[FRONTEND] Cleanup nv_tma_desc_type (triton-lang#6508)

This is now dead code as `nv_tma_desc` only exists in the backend.

Author: peterbell10

python/triton/compiler/code_generator.py

@@ -11,7 +11,7 @@
 from .. import language
 from .._C.libtriton import ir
 from ..language import constexpr, semantic, str_to_ty, tensor
-from ..language.core import _unwrap_if_constexpr, nv_tma_desc_type, base_value, base_type
+from ..language.core import _unwrap_if_constexpr, base_value, base_type
 from ..runtime.jit import get_jit_fn_file_line
 # ideally we wouldn't need any runtime component
 from ..runtime import JITFunction
@@ -258,10 +258,6 @@ def make_template(ty):
             for attr_name, attr_val in attr_specs:
                 if attr_path in val_paths:
                     fn.set_arg_attr(val_paths.index(attr_path), attr_name, attr_val)
-        for i, path in enumerate(val_paths):
-            ty = get_iterable_path(self.arg_types, path)
-            if isinstance(ty, nv_tma_desc_type):
-                fn.set_arg_attr(i, "tt.nv_tma_desc", 1)
         # > add IR values to the template
         cursor = 0
         handles = [fn.args(i) for i in range(fn.get_num_args())]

python/triton/language/__init__.py

@@ -90,7 +90,6 @@
     permute,
     pi32_t,
     pointer_type,
-    nv_tma_desc_type,
     program_id,
     range,
     reduce,
@@ -219,7 +218,6 @@
     "philox_impl",
     "pi32_t",
     "pointer_type",
-    "nv_tma_desc_type",
     "program_id",
     "rand",
     "rand4x",
@@ -293,9 +291,6 @@ def str_to_ty(name):
         block = block_type(dtype, block_shape)
         return tensor_descriptor_type(block, shape_type, stride_type)
 
-    if name == "nvTmaDesc":
-        return nv_tma_desc_type()
-
     if name == "constexpr":
         return constexpr
 

python/triton/language/core.py

@@ -643,13 +643,6 @@ def mangle(self) -> str:
         return f"P{self.element_ty.mangle()}"
 
 
-class nv_tma_desc_type(pointer_type):
-
-    def __init__(self, const=True, address_space=0):
-        super().__init__(uint8, const=const, address_space=address_space)
-        self.name = 'nv_tma_desc_type'
-
-
 class block_type(dtype):
 
     def __init__(self, element_ty: dtype, shape: List):

python/tutorials/06-fused-attention.py

@@ -15,11 +15,16 @@
 
 import pytest
 import torch
-from triton.tools.tensor_descriptor import TensorDescriptor
 
 import triton
 import triton.language as tl
 
+try:
+    from triton.tools.tensor_descriptor import TensorDescriptor
+    HAS_TENSOR_DESC = True
+except ModuleNotFoundError:
+    HAS_TENSOR_DESC = False
+
 DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
@@ -32,15 +37,7 @@ def is_cuda():
 
 
 def supports_tma():
-    return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
-
-
-HAS_TMA_DESC = "nv_tma_desc_type" in dir(tl)
-
-if HAS_TMA_DESC:
-    print("TMA benchmarks will be running with experimental grid constant TMA descriptor.", )
-else:
-    print("TMA benchmarks will be running without grid constant TMA descriptor.", )
+    return HAS_TENSOR_DESC and is_cuda() and torch.cuda.get_device_capability()[0] >= 9
 
 
 @triton.jit