Could nsight compute really trace the OpenAI-triton code?

Developer Tools Nsight Compute
1

As the @triton.jit kernels have been compiled, can the source and SASS catch the relationship between metrics like (live regs) (stall samples)?

image
image2386×1246 154 KB

2

Hi, @pannenets.f

Sorry I don’t quite get your question. What exactly you want to know ? Which kind of relationship do you refer to ?

3

The source and SASS panel show the color bars along with the Python code.

e.g.

@triton.jit
def foo(
  x_ptr,
  y_ptr,
):
  cur_loc = tl.program_id(0)      # color bar 1
  xx = tl.load(x_ptr + cur_loc)   # color bar 2
  yy = xx * xx + xx + 1           # color bar 3
  tl.store(y_ptr + cur_loc, yy).  # color bar 4

Could the color bars reflect things like the register pressure of corresponding lines? As the triton is compiled before running, I wonder the accuracy of tracing.

Thanks! @veraj

4

Yes the Source Page in Nsight Compute will show metrics at the source and SASS instruction level.

Refer the Nsight Compute document Source page section.

5

Hi, @pannenets.f

Sorry for the late response. Have checked with internal dev, the data is showed to python source to the best that the line table information can provide.