search

triton-lang / triton

Development repository for the Triton language and compiler
https://triton-lang.org/
MIT License
13.54k stars 1.67k forks source link

[IR] Add typing for tensor descriptor types #5147

Closed peterbell10 closed 2 weeks ago

peterbell10 commented 2 weeks ago
#### Commits in this PR 1. [IR] Add typing for tensor descriptor types Currently tensor descriptors are just typed as `!tt.ptr` which is exposing the assumption it's using a TMA descriptor. This changes it to a custom type `!tt.tensordesc>` which is lowered to a pointer type in the LLVM IR. I also add two new IR Ops which are used to cast between pointers and tensordesc objects. ```mlir tt.reinterpret_tensor_descriptor %ptr : !tt.ptr to !tt.tensordesc<...> triton_nvidia_gpu.tensor_desc_to_tma_ptr %desc : !tt.tensordesc<...> -> !tt.ptr ``` Really both of these should be nvidia-specific but the first is exposed in the triton IR to keep support for the by-value TMA descriptor API around while we figure out if it's possible to update to the new style. 1. canonicalize -> fold 1. Revert 'canonicalize -> fold' #### [PR chain](https://github.com/jlebar/git-pr-chain) 1. 👉 #5147 👈 **YOU ARE HERE**