[PYTHON][TOOLS] simplify construction of tensor descriptor (#7159)

base dtype may not be (and doesn't need to be) canonicalizable

Author: ptillet

python/triton/language/core.py

@@ -409,6 +409,7 @@ def __init__(self, name):
         self.name = name
         assert name in dtype.SINT_TYPES + dtype.UINT_TYPES + dtype.FP_TYPES + dtype.OTHER_TYPES, name
         self.primitive_bitwidth = get_primitive_bitwidth(name)
+        self.itemsize = self.primitive_bitwidth // 8
         if name in dtype.SINT_TYPES:
             self.int_signedness = dtype.SIGNEDNESS.SIGNED
             self.int_bitwidth = self.primitive_bitwidth

python/triton/tools/tensor_descriptor.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from typing import List, Any
-from triton._utils import validate_block_shape, canonicalize_dtype, get_primitive_bitwidth
+from triton._utils import validate_block_shape
 
 
 @dataclass
@@ -18,8 +18,7 @@ def __post_init__(self):
         assert rank <= 5, "rank cannot be more than 5"
         assert self.base.data_ptr() % 16 == 0, "base must be 16-byte aligned"
         validate_block_shape(self.block_shape)
-        dtype_str = canonicalize_dtype(self.base.dtype)
-        elem_bytes = get_primitive_bitwidth(dtype_str) // 8
+        elem_bytes = self.base.dtype.itemsize
         for stride in self.strides[:-1]:
             assert (stride * elem_bytes) % 16 == 0, "strides must be 16-byte aligned"
         assert self.strides[-1] == 1, "Last dimension must be contiguous"