[Gluon] Use _convert_elem_to_ir_value in some APIs (#7022)

This allows passing constexprs as the arguments, like `tmem.subslice(0)`

Author: Mogball

python/test/gluon/test_frontend.py

@@ -541,6 +541,29 @@ def kernel():
     assert "order must be a permutation of 0..(rank-1), but was [1]" in str(e.value.__cause__)
 
 
+@gluon.jit
+def tmem_subslice_kernel():
+    layout: ttgl.constexpr = ttgl.nvidia.blackwell.TensorMemoryLayout(block=[128, 128], unpacked=True)
+    tmem = ttgl.nvidia.blackwell.allocate_tensor_memory(ttgl.int32, [2, 256, 256], layout)
+    tmem.subslice(0)
+
+
+def test_tmem_subslice_constexpr():
+    expecttest.assert_expected_inline(
+        run_parser(tmem_subslice_kernel).str_nodebug(), """\
+#tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, unpacked = true>
+module {
+  tt.func public @tmem_subslice_kernel() attributes {noinline = false} {
+    %result = ttng.tmem_alloc : () -> !ttg.memdesc<2x256x256xi32, #tmem, #ttng.tensor_memory, mutable>
+    %c0_i32 = arith.constant 0 : i32
+    %c0_i32_0 = arith.constant 0 : i32
+    %0 = ttg.memdesc_subview %result[%c0_i32, %c0_i32_0, %c0_i32_0] : !ttg.memdesc<2x256x256xi32, #tmem, #ttng.tensor_memory, mutable> -> !ttg.memdesc<256x256xi32, #tmem, #ttng.tensor_memory, mutable, 2x256x256>
+    tt.return
+  }
+}
+""")
+
+
 @gluon.jit
 def smem_and_layout_user(smem, a: ttgl.constexpr):
     pass

python/triton/experimental/gluon/language/_semantic.py

@@ -71,7 +71,7 @@ def memdesc_slice(mem_desc, index, shape, layout, builder: GluonOpBuilder):
     assert mem_desc.rank > len(shape), f"source rank ({mem_desc.rank}) must be greater than result rank ({len(shape)})"
 
     offsets = [builder.get_int32(0)] * mem_desc.rank
-    offsets[0] = index.handle
+    offsets[0] = tl_semantic._convert_elem_to_ir_value(builder, index, require_i64=False)
     return _memdesc_subview(mem_desc, offsets, shape, layout, builder)
 
 

python/triton/experimental/gluon/language/nvidia/blackwell/__init__.py

@@ -2,6 +2,7 @@
 from typing import Optional, Tuple, List, TYPE_CHECKING
 
 from dataclasses import dataclass
+from triton.language.semantic import _convert_elem_to_ir_value, _convert_to_ir_values
 from triton.experimental.gluon.language import _core as ttgl
 from triton.experimental.gluon.language._core import builtin, base_type, base_value, _unwrap_if_constexpr
 
@@ -140,12 +141,12 @@ def subslice(self, index, shape=None, layout=None, _builder: GluonOpBuilder = No
         if shape is None:
             shape = self.shape[1:]
 
-        index = _unwrap_if_constexpr(index)
+        index = _convert_elem_to_ir_value(_builder, index, require_i64=False)
         shape = [_unwrap_if_constexpr(s) for s in shape]
         layout = _unwrap_if_constexpr(layout)
 
         offsets = [_builder.get_int32(0)] * self.rank
-        offsets[0] = index.handle
+        offsets[0] = index
         ret = tensor_memory_descriptor(None, self.dtype, shape, layout, self.type.alloc_shape)
         ret.handle = _builder.create_memdesc_subview(ret.type.to_ir(_builder), self.handle, offsets)
         return ret
@@ -178,6 +179,6 @@ def tcgen05_mma(a, b, acc, *, use_acc=True, pred=True, mbarriers=None, mbarrier_
             true = ttgl.to_tensor(True, _builder=_builder)
             mbarrier_preds = [true] * len(mbarriers)
         else:
-            mbarrier_preds = [pred.handle for pred in mbarrier_preds]
+            mbarrier_preds = _convert_to_ir_values(_builder, mbarrier_preds, require_i64=False)
 
     _builder.create_tcgen05_mma(a.handle, b.handle, acc.handle, use_acc.handle, pred.handle, mbarriers, mbarrier_preds)

python/triton/language/semantic.py

@@ -1865,6 +1865,8 @@ def _convert_elem_to_ir_value(builder, elem, require_i64):
     if isinstance(elem, int):
         elem = tl.constexpr(elem)
     if isinstance(elem, tl.constexpr):
+        if isinstance(elem.value, bool):
+            return builder.get_int1(elem.value)
         if require_i64:
             assert -2**63 <= elem.value < 2**63, f"Block pointers only support 64 bit `shape/strides`, " \
                 f"got a value {elem.value} which is out of the range"