Add update_tensor_descriptor operation to Triton/Gluon

[alexsamardzic]

This PR is to add update_tensor_descriptor operation, to simplify writing kernels like grouped MM (like pytorch/pytorch#166063, in particular to avoid handling special cases like this). This is also to match update_tensormap op in CuTe DSL, like used here.

The operation reads the existing descriptor from GMEM into SMEM, performs updates in SMEM, and writes updated descriptor back into GMEM. So the rationale for using this operation instead of creating a new descriptor is to save a GMEM allocation (more precisely, to trade it for reading from GMEM), and to emit only tensormap.replace.tile.* PTX instructions for the descriptor fields that are actually changed. Otherwise, the implementation closely follows make_tensor_descriptor implementation. The end-to-end performance improvement is minor, but the main advantage is that the code for the cases like the kernel pointed above is cleaner. This PR makes it possible to change tensor base pointer, shape and strides fields in the descriptor; changing other fields could be added in the future if there is a need.

[peterbell10]

The in-place update kind of forces the underlying implementation to be memory-backed which may not be the case for all hardware, including pre-hopper where we translate tensor descriptors to normal pointer indexing.

[alexsamardzic]

There are several options to fix it, for example lowering basically to what make_tensor_descriptor does for older hardware. But overall: I'm actually not sure having this operation in Triton is appropriate, so how about removing the Triton version, and moving Gluon version under tma or better hopper namespace?

[alexsamardzic]

Added a commit that removes Triton version of the operation from PR, and moves Gluon version under tma namespace.

[peterbell10]

I think this is illegal. We're passing the descriptor in param space which should be constant.

I also think this will break if your launch grid is larger than the number of SMs. I expect the second program to be scheduled on a single SM would see the already updated descriptor.

[alexsamardzic]

Uhm, indeed. Would it be acceptable to limit the operation to work only for the descriptors created from within the kernel, and thus avoid both problems you pointed to?

[alexsamardzic]

Added another commit to implement the change proposed.