@zou3519 @bdhirsh

This is surprising. It fails with

RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

in P857566445.

Looking at it with a debugger, we are calling copy_ on a faketensor that requires grad. Looks like aot_eager does not execute this codepath.

@jansel It looks like all the existing kernels are CachingAutotuner but in my case I end up with JitFunction from triton. I assume in order to support @triton.autotune I'm gonna need to make this work with cachingautotuner but for the time being is this fine?

We should not use JitFunction from Triton, we want parallel ahead of time compiles.

You should be able to do something similar to:

And put a @template above the generated Triton kernel. You will need to generate proper meta.

Clones are removed as part of remove_noop_ops. Looking at the output code, I think the clones are still there, so i want to check whether to run remove_noop_ops again or move this check above

I think it's better to do this as part of reinplace_scatters.

1. The clone removal pass doesn't handle inplace mutations.
2. reinplace_scatters is exactly for this purpose - it converts scatter (i.e. clone + scatter_) into just scatter_. It also converts scatter + copy_ into just scatter_ as well.

oh I see - is inductor's clone removal pass done during lowering time, so if we desugar the functional triton_kernel op into mutable op + clones, it will be "too late"?

hmm... maybe a little late, but @Chillee I don't think that inductor's fx graph caching TORCHINDUCTOR_FX_GRAPH_CACHE=1 will play well with this kernel_side_table.

What do you think - should we just make sure not to add to the cache if our graph uses any user triton kernels, to unblock, or fully revisit the side table?

It seems to me like the FX graph that we get (both from dynamo and AOTAutograd) will burn in a random index that's supposed to map to a user triton kernel that dynamo saw:

def forward(self, x_1, output_1):
    triton_kernel_wrapper_functional_proxy = torch._higher_order_ops.triton_kernel_wrap.triton_kernel_wrapper_functional(kernel_idx = 0, grid = (5,), kwargs = {'in_ptr0': x_1, 'out_ptr': output_1, 'n_elements': 5, 'BLOCK_SIZE': 16});  x_1 = #
    getitem = triton_kernel_wrapper_functional_proxy['in_ptr0']
    getitem_1 = triton_kernel_wrapper_functional_proxy['out_ptr']
    getitem_2 = triton_kernel_wrapper_functional_proxy['n_elements']
    getitem_3 = triton_kernel_wrapper_functional_proxy['BLOCK_SIZE'];  triton_kernel_wrapper_functional_proxy = None
    return getitem_1

But if we were to map this graph to a cached, compiled inductor graph, we have no guarantee that that index will map to the same triton kernel.

Maybe we can have the cache key also include the kernel_side_table's mapping, from index to triton kernel (or a hash of its source code)?

Yeah this is why I was suggesting that it might be better to just include the Triton source code in the operator definition.

We would not only need to store the source code of one function but also every single dependancy. Another downside of that is the FX graph would be impossible to read.

If this is the preferred solution, I'm happy to change it but IMO caching kernel_side_table's mapping seems like a cleaner solution.

We should not use JitFunction from Triton, we want parallel ahead of time compiles.

You should be able to do something similar to:

And put a @template above the generated Triton kernel. You will need to generate proper meta.