nit: V.graph.is_inference already exists, see:

Could we just add is_backward to GraphLowering object and query that for the is_backward property, instead of adding this ?

Totally, I missed the is_inference property. Will add is_backward there as well and revert this commit!

Reverted this commit and added is_backward to GraphLowering 6a27b68 720b971

? add full docstring ?

This depends on whether RangeHandle approach from above is the right way to go. I'll update the doc to cover the current implementation.

I dont know a ton about nvtx.. not maybe @ezyang or @eqy might be able to weigh in

Added some docs 6390c08
Please, let me know if I can make it clearer 🙇

cc @ezyang, I see that you first added nvrtx bindings (7 years ago). do you want to take a look ?

can you add this conditionally ? see

@eellison done a61c703

would you mind posting an example output code for a fused kernel ? this would also be a good candidate for a test see run_and_get_code

Added a couple of test cases here 7d623ef
Example output is in the comment below: #130429 (comment)

This is pretty cuda specific codegen for the general scheduler.. also, I wonder if any code here could be shared with the pytorch profiler. cc @davidberard98

for profiler, we put the code in the runtime so we don't fill the codegen with profiling annotations - is this viable for your use case?

I'm not sure how to make it less cuda-specific unfortunately though. @sraikund16 tells me that the NVTX handling in profiler is somewhat cuda-specific

for profiler, we put the code in the runtime so we don't fill the codegen with profiling annotations - is this viable for your use case?

It seems the profiler works at a more fine grained level? i.e. it "wraps" kernel runs into profiling events? These annotations work at a slightly higher level/granularity. I think the training annotations (bw/fw/inference) can be moved outside of the codegen, but I'm not sure about the "buffer" annotations 🤔

Re: CUDA specific: this is indeed the case, not sure how to handle this properly. I can think of emitting some "abstract" Begin/EndAnnotationLine with nvtx calls hidden there, but I'm not certain if it brings much value? Happy to find the right solution 🙌

The buffer annotation, like the profiler, is wrapping a single run of a kernel.

        buffer_annotation = nvtx.device_range_start('op1')
        buf1 = empty_strided_cuda((5, ), (1, ), torch.float32)
        # Topologically Sorted Source Nodes: [mul], Original ATen: [aten.mul]
        triton_poi_fused_mul_1.run(arg0_1, arg1_1, buf1, 5, grid=grid(5), stream=stream0)
        del arg0_1
        del arg1_1
        nvtx.device_range_end(buffer_annotation)

I think it would be an improvement to put the logic here in the same place as the profiler.

As for as CUDA specific - I commented elsewhere about moving the codegen logic elsewhere.

I agree that it would be an improvement, but it would change the granularity of the annotations a bit: current "buffer annotations" cover kernel run and all the memory allocations/deallocations, i.e. "annotate everything that happens to compute the buffer." Wrapping kernel runs would be more of an "annotate kernels" which is somewhat orthogonal to the current version. Perhaps it could be a third annotation option? 🤔

TBH, i'm not convinced that this is really that meaningful. It's not tracking memory (i.e, dont know when a particular buffer is allocated/deallocated), and the memory allocations/deallocations themselves are extremely cheap and happen on a different timeline than cuda because cuda is async. In any case, lets move this out of scheduler.py if you are convinced on keeping buffer or put in profiler

In any case, lets move this out of scheduler.py if you are convinced on keeping buffer or put in profiler

Apologies for the back and forth, just to ensure I understand this correctly: moving this into profiler implies that the models must use the profiler and the annotations would only appear in case the profiling is enabled?
if that's the case I'd prefer to not move it into profiler as it makes the integration into an arbitrary model harder.

Additionally, there are two more concerns:

• it doesn't seem like the buffer names are available around the run method? So the best we can do is to add annotations based on the kernel name, but the same kernel can be used to compute several distinct buffers. It's of course possible to pass the buffer names around, but it doesn't look like a particularly good idea?
• the profiler is called from within triton, which would miss non-triton kernels in case of mixed execution environment

For moving it into wrapper.py, I guess it'd require wrapping all the kernel calls into special lines (e.g. KernelCallLine) and adding another check here? I don't see any special *Line classes for such invocations (the last time I checked they were simply Python strings at the wrapper level). The buffer names are also missing at this level, though.

Does it make sense? Or I'm missing something?

@eellison I moved training annotations into wrapper.py 78c1c7e

Regarding the buffer annotations: I cannot find a better place to emit this code than within Scheduler's codegen.
It's possible to make it more general by introducing special, generic *Lines and moving the actual emission (together with the config check) into wrapper.py, but that would still leave "traces" in the scheduler.

If adding this to Scheduler is a strong no-go, then I could remove it from this PR leaving only "training annotations" here. I'd be happy to open a followup PR with buffer/kernel annotations for further discussion.

Let's try to keep codegen aspects out of scheduler and keep high level scheduler logic lean. This annotation can go here:

Similarly, the end can go in _generate.

The buffer annotation, like the profiler, is wrapping a single run of a kernel.

        buffer_annotation = nvtx.device_range_start('op1')
        buf1 = empty_strided_cuda((5, ), (1, ), torch.float32)
        # Topologically Sorted Source Nodes: [mul], Original ATen: [aten.mul]
        triton_poi_fused_mul_1.run(arg0_1, arg1_1, buf1, 5, grid=grid(5), stream=stream0)
        del arg0_1
        del arg1_1
        nvtx.device_range_end(buffer_annotation)

I think it would be an improvement to put the logic here in the same place as the profiler.

As for as CUDA specific - I commented elsewhere about moving the codegen logic elsewhere.