Merge commit 'e19644610878ab19ec375d3eabfb464398aade32'

Author: whitneywhtsang

README.md

@@ -245,6 +245,7 @@ For detailed instructions on how to debug Triton's frontend, please refer to thi
 - `TRITON_KERNEL_OVERRIDE` enables the override of the compiled kernel with a user-specified IR/ptx/amdgcn at the beginning of each compilation stage.
 - `TRITON_OVERRIDE_DIR` specifies the directory from which to load the IR/ptx/amdgcn files when `TRITON_KERNEL_OVERRIDE` is set to 1.
 - `TRITON_F32_DEFAULT` sets the default input precision of `tl.dot` when using 32-bit floats, which can be either `ieee`, `tf32`, or `tf32x3`.
+- `TRITON_FRONT_END_DEBUGGING=1` disables exception wrapping when an error occurs in the compiler frontend, allowing the full stack trace to be seen.
 
 **Kernel Override Steps**
 

bin/RegisterTritonDialects.h

@@ -85,7 +85,6 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   // TritonAMDGPUToLLVM passes
   mlir::triton::registerConvertTritonAMDGPUToLLVM();
   mlir::triton::registerConvertBuiltinFuncToLLVM();
-  mlir::triton::registerDecomposeUnsupportedAMDConversions();
   mlir::triton::registerOptimizeAMDLDSUsage();
 
   // TritonAMDGPUTransforms passes

include/triton/Conversion/TritonGPUToLLVM/Patterns.h

@@ -6,7 +6,6 @@ add_triton_library(TritonGPUToLLVM
     AssertOpToLLVM.cpp
     ControlFlowOpToLLVM.cpp
     ConvertLayoutOpToLLVM.cpp
-    DecomposeUnsupportedConversions.cpp
     ElementwiseOpToLLVM.cpp
     FuncOpToLLVM.cpp
     GatherOpToLLVM.cpp

lib/Conversion/TritonGPUToLLVM/CMakeLists.txt

@@ -421,6 +421,41 @@ def alloc_fn(size: int, align: int, stream: Optional[int]) -> torch.Tensor:
     torch.testing.assert_close(expect, out)
 
 
+@triton.jit(noinline=True)
+def tensor_descriptor_arg_helper(in_desc, out_desc, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr):
+    moffset = tl.program_id(0) * M_BLOCK
+    noffset = tl.program_id(1) * N_BLOCK
+    value = in_desc.load([moffset, noffset])
+    out_desc.store([moffset, noffset], value.abs())
+
+
+@requires_tma
+@pytest.mark.interpreter
+def test_tensor_descriptor_argument():
+
+    @triton.jit
+    def kernel(out_ptr, a_ptr, M, N, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr):
+        out_desc = tl.make_tensor_descriptor(out_ptr, shape=[M, N], strides=[N, 1], block_shape=[M_BLOCK, N_BLOCK])
+        in_desc = tl.make_tensor_descriptor(a_ptr, shape=[M, N], strides=[N, 1], block_shape=[M_BLOCK, N_BLOCK])
+        tensor_descriptor_arg_helper(in_desc, out_desc, M_BLOCK, N_BLOCK)
+
+    M, N = 32, 128
+    inp = torch.randn((M, N), device="cuda")
+
+    M_BLOCK = 8
+    N_BLOCK = 32
+    out = inp.new_zeros((M, N))
+
+    def alloc_fn(size: int, align: int, stream: Optional[int]) -> torch.Tensor:
+        return torch.empty(size, dtype=torch.int8, device="cuda")
+
+    triton.set_allocator(alloc_fn)
+
+    expect = inp.abs()
+    kernel[(M // M_BLOCK, N // N_BLOCK)](out, inp, M, N, M_BLOCK, N_BLOCK)
+    torch.testing.assert_close(expect, out)
+
+
 @triton.jit
 def matmul_kernel_make_tensor_desciptor(a_ptr, b_ptr, c_ptr,  #
                                         M, N, K,  #

lib/Conversion/TritonGPUToLLVM/DecomposeUnsupportedConversions.cpp

@@ -16,6 +16,7 @@
 # ideally we wouldn't need any runtime component
 from ..runtime import JITFunction
 from .._utils import find_paths_if, get_iterable_path, set_iterable_path
+from . import config
 
 from .errors import (CompilationError, CompileTimeAssertionFailure, UnsupportedLanguageConstruct)
 
@@ -27,29 +28,9 @@ def check_identifier_legality(name, type):
     return name
 
 
-def mangle_ty(ty):
-    if ty.is_tuple():
-        return 'T' + '_'.join(map(mangle_ty, ty.types)) + 'T'
-    if ty.is_ptr():
-        return 'P' + mangle_ty(ty.element_ty)
-    if ty.is_int():
-        SIGNED = language.dtype.SIGNEDNESS.SIGNED
-        prefix = 'i' if ty.int_signedness == SIGNED else 'u'
-        return prefix + str(ty.int_bitwidth)
-    if ty.is_floating():
-        return str(ty)
-    if ty.is_block():
-        elt = mangle_ty(ty.scalar)
-        shape = '_'.join(map(str, ty.shape))
-        return f'{elt}S{shape}S'
-    if ty.is_void():
-        return 'V'
-    raise TypeError(f'Unsupported type {ty}')
-
-
 def mangle_fn(name, arg_tys, constants):
     # doesn't mangle ret type, which must be a function of arg tys
-    mangled_arg_names = '_'.join([mangle_ty(ty) for ty in arg_tys])
+    mangled_arg_names = '_'.join([ty.mangle() for ty in arg_tys])
     mangled_constants = '_'.join([f'{i}c{repr(constants[i])}' for i in sorted(constants)])
     mangled_constants = mangled_constants.replace('.', '_d_')
     mangled_constants = mangled_constants.replace("'", '_sq_')
@@ -71,8 +52,8 @@ def _is_constexpr(o: Any) -> bool:
     return o is None or isinstance(o, (constexpr, language.core.dtype))
 
 
-def _is_triton_scalar(o: Any) -> bool:
-    return _is_triton_tensor(o) and (not o.type.is_block() or o.type.numel == 1)
+def _is_non_scalar_tensor(o: Any) -> bool:
+    return _is_triton_tensor(o) and (o.type.is_block() and o.type.numel != 1)
 
 
 def _is_list_like(o: Any) -> bool:
@@ -82,7 +63,7 @@ def _is_list_like(o: Any) -> bool:
 def _check_fn_args(node, fn, args):
     if fn.noinline:
         for idx, arg in enumerate(args):
-            if not _is_constexpr(arg) and not _is_triton_scalar(arg):
+            if not _is_constexpr(arg) and _is_non_scalar_tensor(arg):
                 raise UnsupportedLanguageConstruct(
                     fn.src, node,
                     f'Function {fn.__name__} is marked noinline, but was called with non-scalar argument {fn.arg_names[idx]}:{arg}'
@@ -241,26 +222,26 @@ def __init__(self, ret_types, arg_types, constants, attrs):
         self.constants = constants
         self.attrs = attrs
 
-    def return_types_ir(self, builder: ir.builder):
-        ret_types = []
-        for ret_ty in self.ret_types:
-            if ret_ty is None:
+    def flatten_ir_types(self, builder: ir.builder, types: List[base_type]) -> List[ir.type]:
+        ir_types = []
+        for ty in types:
+            if ty is None:
                 continue
-            ir_ty = ret_ty.to_ir(builder)
-            if isinstance(ir_ty, list):
-                ret_types.extend(ir_ty)
-            else:
-                ret_types.append(ir_ty)
-        return ret_types
+            ty._flatten_ir_types(builder, ir_types)
+        return ir_types
+
+    def return_types_ir(self, builder: ir.builder) -> List[ir.type]:
+        return self.flatten_ir_types(builder, self.ret_types)
 
     def serialize(self, builder: ir.builder):
         # fill up IR values in template
         # > build function
         is_val = lambda path, _: path not in self.constants and _ is not None
         val_paths = list(find_paths_if(self.arg_types, is_val))
-        arg_types = [get_iterable_path(self.arg_types, path).to_ir(builder) for path in val_paths]
-        ret_types = self.return_types_ir(builder)
-        return builder.get_function_ty(arg_types, ret_types)
+        arg_types = [get_iterable_path(self.arg_types, path) for path in val_paths]
+        arg_types_ir = self.flatten_ir_types(builder, arg_types)
+        ret_types_ir = self.return_types_ir(builder)
+        return builder.get_function_ty(arg_types_ir, ret_types_ir)
 
     def deserialize(self, fn):
         # create "template"
@@ -282,9 +263,12 @@ def make_template(ty):
             if isinstance(ty, nv_tma_desc_type):
                 fn.set_arg_attr(i, "tt.nv_tma_desc", 1)
         # > add IR values to the template
-        for i, path in enumerate(val_paths):
+        cursor = 0
+        handles = [fn.args(i) for i in range(fn.get_num_args())]
+        for path in val_paths:
             ty = get_iterable_path(self.arg_types, path)
-            set_iterable_path(vals, path, language.tensor(fn.args(i), ty))
+            val, cursor = ty._unflatten_ir(handles, cursor)
+            set_iterable_path(vals, path, val)
         # > add constexpr values to the template
         constants = self.constants
         for path, val in constants.items():
@@ -1218,14 +1202,16 @@ def call_JitFunction(self, fn: JITFunction, args, kwargs):
                 generator.visit(fn.parse())
             except Exception as e:
                 # Wrap the error in the callee with the location of the call.
+                if config.front_end_debugging():
+                    raise
                 raise CompilationError(self.jit_fn.src, self.cur_node, None) from e
 
             callee_ret_type = generator.ret_type
             self.function_ret_types[fn_name] = callee_ret_type
         else:
             callee_ret_type = self.function_ret_types[fn_name]
         symbol = self.module.get_function(fn_name)
-        args_val = [arg.handle for arg in args_val]
+        args_val = flatten_values_to_ir(args_val)
         call_op = self.builder.call(symbol, args_val)
         if callee_ret_type == language.void:
             return None
@@ -1256,6 +1242,8 @@ def visit_Call(self, node):
                     ret = language.tuple(ret)
                 return ret
             except Exception as e:
+                if config.front_end_debugging():
+                    raise
                 # Normally when we raise a CompilationError, we raise it as
                 # `from None`, because the original fileline from the exception
                 # is not relevant (and often points into code_generator.py
@@ -1335,6 +1323,8 @@ def visit(self, node):
             except CompilationError:
                 raise
             except Exception as e:
+                if config.front_end_debugging():
+                    raise
                 # Wrap the error in a CompilationError which contains the source
                 # of the @jit function.
                 raise CompilationError(self.jit_fn.src, self.cur_node, repr(e)) from None

python/test/unit/cuda/test_tensor_descriptor.py

@@ -11,6 +11,7 @@
 from ..tools.disasm import get_sass, get_spvdis
 # TODO: this shouldn't be here
 from .code_generator import ast_to_ttir
+from . import config
 from pathlib import Path
 import re
 import functools
@@ -181,7 +182,7 @@ def filter_traceback(e: BaseException):
 
     These are uninteresting to the user -- "just show me *my* code!"
     """
-    if os.getenv("TRITON_FRONT_END_DEBUGGING", "0") == "1":
+    if config.front_end_debugging():
         return
 
     if e.__cause__ is not None:

python/triton/compiler/code_generator.py

@@ -0,0 +1,5 @@
+import os
+
+
+def front_end_debugging():
+    return os.getenv("TRITON_FRONT_END_DEBUGGING", "0") == "1"