Add alias support for tle.gpu.alloc to enable shared memory aliasing

lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp

@@ -38,6 +38,7 @@
 #include "triton/Analysis/Utility.h"
 #include "triton/Conversion/MLIRTypes.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
+#include "tle/dialect/include/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/OpInterfaces.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Attributes.h"
@@ -92,6 +93,10 @@ static Value getUnderlyingMemDesc(Value value) {
       value = subslice.getSrc();
       continue;
     }
+    if (auto alias = value.getDefiningOp<tle::MemDescAliasOp>()) {
+      value = alias.getSrc();
+      continue;
+    }
     if (auto trans = value.getDefiningOp<MemDescTransOp>()) {
       value = trans.getSrc();
       continue;

lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp

@@ -84,6 +84,10 @@ static Value stripMemDescViews(Value value) {
       value = subslice.getSrc();
       continue;
     }
+    if (auto alias = value.getDefiningOp<tle::MemDescAliasOp>()) {
+      value = alias.getSrc();
+      continue;
+    }
     if (auto trans = value.getDefiningOp<MemDescTransOp>()) {
       value = trans.getSrc();
       continue;

lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.cpp

@@ -359,6 +359,10 @@ static Value stripMemDescViewsForLocalAlloc(Value value) {
       value = subslice.getSrc();
       continue;
     }
+    if (auto alias = value.getDefiningOp<tle::MemDescAliasOp>()) {
+      value = alias.getSrc();
+      continue;
+    }
     if (auto trans = value.getDefiningOp<ttg::MemDescTransOp>()) {
       value = trans.getSrc();
       continue;
@@ -548,6 +552,10 @@ static MemDescResource getMemDescResource(Value value) {
       value = subslice.getSrc();
       continue;
     }
+    if (auto alias = value.getDefiningOp<tle::MemDescAliasOp>()) {
+      value = alias.getSrc();
+      continue;
+    }
     if (auto trans = value.getDefiningOp<ttg::MemDescTransOp>()) {
       value = trans.getSrc();
       continue;

lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp

@@ -37,6 +37,7 @@
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Attributes.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "tle/dialect/include/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/PipelineExpander.h"
 #include "triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h"
 #include "triton/Dialect/TritonGPU/Transforms/Schedule.h"
@@ -479,7 +480,8 @@ static std::optional<int> dotCanBeProperlyAsync(ttng::WarpGroupDotOp dotOp,
     // allowed in between.
     Value transitiveOperand = operand;
     while (isa_and_nonnull<ttg::ConvertLayoutOp, ttg::MemDescTransOp,
-                           ttg::MemDescReshapeOp, ttg::MemDescSubsliceOp>(
+                           ttg::MemDescReshapeOp, ttg::MemDescSubsliceOp,
+                           mlir::triton::tle::MemDescAliasOp>(
                transitiveOperand.getDefiningOp()) ||
            isa<BlockArgument>(transitiveOperand)) {
       auto blockArg = dyn_cast<BlockArgument>(transitiveOperand);

lib/Dialect/TritonGPU/Transforms/Prefetch.cpp

@@ -60,6 +60,7 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "tle/dialect/include/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/Passes.h"
 #include "llvm/Support/Debug.h"
 

lib/Dialect/TritonGPU/Transforms/Utility.cpp

@@ -1494,7 +1494,8 @@ ttg::LocalAllocOp findShmemAlloc(Value operand) {
   // allowed in between.
   Value transitiveOperand = operand;
   while (isa_and_nonnull<ttg::ConvertLayoutOp, tt::TransOp, ttg::MemDescTransOp,
-                         ttg::MemDescReshapeOp, ttg::MemDescSubsliceOp>(
+                         ttg::MemDescReshapeOp, ttg::MemDescSubsliceOp,
+                         tle::MemDescAliasOp>(
              transitiveOperand.getDefiningOp()) ||
          isa<BlockArgument>(transitiveOperand)) {
     if (auto blockArg = dyn_cast<BlockArgument>(transitiveOperand)) {
@@ -1689,6 +1690,14 @@ void replaceUsesAndPropagateType(
           oldType.getMemorySpace(), isMutable, oldType.getAllocShape());
       newVal = ttg::MemDescSubsliceOp::create(
           builder, subslice.getLoc(), newDstType, val, subslice.getOffsets());
+    } else if (auto alias = dyn_cast<tle::MemDescAliasOp>(user)) {
+      ttg::MemDescType oldType = alias.getType();
+      bool isMutable = cast<ttg::MemDescType>(val.getType()).getMutableMemory();
+      Type newDstType = ttg::MemDescType::get(
+          oldType.getShape(), oldType.getElementType(), oldType.getEncoding(),
+          oldType.getMemorySpace(), isMutable, oldType.getAllocShape());
+      newVal = tle::MemDescAliasOp::create(
+          builder, alias.getLoc(), newDstType, val, alias.getOffsetBytesAttr());
     } else if (auto trans = dyn_cast<ttg::MemDescTransOp>(user)) {
       newVal = ttg::MemDescTransOp::create(builder, trans.getLoc(), val,
                                            trans.getOrder());

python/test/tle/unit/test_tle_alloc_alias.py

@@ -0,0 +1,130 @@
+# flagtree tle
+"""TLE alloc alias unit tests — mock-based, no GPU required."""
+
+import pytest
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+
+class TestAllocAlias:
+
+    class _FakeTensor:
+        def __init__(self, handle, ty):
+            self.handle = handle
+            self.type = ty
+
+    class _FakeBuilder:
+
+        def __init__(self):
+            self.memdesc_type_args = None
+            self.memdesc_alias_args = None
+            self.swizzled_encoding_args = None
+
+        def get_half_ty(self):
+            return "fp16"
+
+        def make_swizzled_shared_encoding_attr(self, vector_size, per_phase, max_phase, order, ctas_per_cga,
+                                               cta_split_num, cta_order):
+            self.swizzled_encoding_args = (
+                vector_size, per_phase, max_phase,
+                list(order), list(ctas_per_cga), list(cta_split_num), list(cta_order),
+            )
+            return "fake_layout"
+
+        def get_memdesc_type(self, shape, element_ty, layout, space, alloc_shape=None):
+            self.memdesc_type_args = (list(shape), element_ty, layout, space, alloc_shape)
+            return ("memdesc", tuple(shape), element_ty, layout, space,
+                    None if alloc_shape is None else tuple(alloc_shape))
+
+        def create_local_alloc(self, *args):
+            return "alloc_handle"
+
+        def create_memdesc_alias(self, result_ty, src, offset_bytes):
+            self.memdesc_alias_args = (result_ty, src, offset_bytes)
+            return "alias_handle"
+
+    class _FakeSemantic:
+
+        def __init__(self):
+            self.builder = TestAllocAlias._FakeBuilder()
+
+        def to_tensor(self, value):
+            if isinstance(value, TestAllocAlias._FakeTensor):
+                return value
+            if isinstance(value, bool):
+                return TestAllocAlias._FakeTensor(f"pred_{value}", tl.int1)
+            if isinstance(value, int):
+                return TestAllocAlias._FakeTensor(f"stage_{value}", tl.int32)
+            raise TypeError(f"unsupported fake tensor input: {value!r}")
+
+    def _make_buffer(self, shape):
+        semantic = self._FakeSemantic()
+        layout = tle.gpu.swizzled_shared_layout.make_default(len(shape))
+        return (
+            tle.gpu.buffered_tensor("base", tl.float16, shape, tle.gpu.smem, layout, semantic),
+            semantic,
+        )
+
+    def test_alloc_alias_creates_typed_memdesc_alias_view(self):
+        """alloc(alias=...) returns a typed view without creating a new allocation."""
+        buffer, semantic = self._make_buffer([4, 16, 32])
+        alias = tle.gpu.alloc(
+            (2, 16, 16),
+            tl.float16,
+            layout=buffer.type.layout,
+            alias=buffer,
+            alias_offset_bytes=64,
+            _semantic=semantic,
+        )
+
+        assert isinstance(alias, tle.gpu.buffered_tensor)
+        assert alias.handle == "alias_handle"
+        assert alias.shape == [2, 16, 16]
+        assert alias.dtype == tl.float16
+        assert alias.type.storage is tle.gpu.smem
+        assert semantic.builder.memdesc_alias_args == (
+            ("memdesc", (2, 16, 16), "fp16", "fake_layout", "smem", None),
+            "base",
+            64,
+        )
+
+    def test_alloc_alias_rejects_init_value(self):
+        buffer, semantic = self._make_buffer([4, 16, 32])
+        init = self._FakeTensor("init", tl.float16)
+
+        with pytest.raises(ValueError, match="alias mode cannot be combined"):
+            tle.gpu.alloc(
+                (2, 16, 16),
+                tl.float16,
+                layout=buffer.type.layout,
+                init_value=init,
+                alias=buffer,
+                _semantic=semantic,
+            )
+
+    def test_alloc_alias_rejects_non_smem_buffer(self):
+        """alias source must be a shared-memory buffered_tensor."""
+        semantic = self._FakeSemantic()
+        fake_buffer = self._FakeTensor("tmem_buf", tl.float16)
+
+        with pytest.raises(ValueError, match="tle.buffered_tensor"):
+            tle.gpu.alloc(
+                (2, 16, 16),
+                tl.float16,
+                alias=fake_buffer,
+                _semantic=semantic,
+            )
+
+    def test_alloc_alias_invalid_offset_type(self):
+        """Bytes are a compile-time argument."""
+        buffer, semantic = self._make_buffer([4, 16, 32])
+
+        with pytest.raises(ValueError, match="compile-time integer"):
+            tle.gpu.alloc(
+                (2, 16, 16),
+                tl.float16,
+                layout=buffer.type.layout,
+                alias=buffer,
+                alias_offset_bytes=b"not_an_int",
+                _semantic=semantic,
+            )

python/triton/experimental/tle/language/gpu/core.py

@@ -255,6 +255,8 @@ def alloc(
     layout: Optional[tle.shared_layout] = None,
     scope: tle.scope = tle.smem,
     init_value: Optional[tl.tensor] = None,
+    alias: Optional[tle.buffered_tensor] = None,
+    alias_offset_bytes: int = 0,
     nv_mma_shared_layout=True,
     _semantic=None,
 ) -> tle.buffered_tensor:
@@ -266,6 +268,9 @@ def alloc(
         dtype: Data type
         layout: Memory layout encoding (optional)
         scope: Storage type (default to shared memory)
+        init_value: Optional initial register tensor for a new allocation
+        alias: Optional source shared-memory buffer to alias instead of allocating
+        alias_offset_bytes: Static byte offset from alias source view base
         _semantic: Semantic analyzer (internal use)
 
     Returns:
@@ -289,6 +294,20 @@ def alloc(
     if not isinstance(scope, tle.scope):
         raise ValueError(f"Storage type must be tle.scope, but got {type(scope)}")
 
+    alias = tl._unwrap_if_constexpr(alias)
+    alias_offset_bytes = tl._unwrap_if_constexpr(alias_offset_bytes)
+    if alias is not None:
+        if init_value is not None:
+            raise ValueError("alloc alias mode cannot be combined with init_value")
+        if not isinstance(alias, tle.buffered_tensor):
+            raise ValueError(f"alias must be a tle.buffered_tensor, but got {type(alias)}")
+        if scope is not tle.smem or alias.type.storage is not tle.smem:
+            raise ValueError("alloc alias mode currently supports only smem buffers")
+        if isinstance(alias_offset_bytes, bool) or not isinstance(alias_offset_bytes, int):
+            raise ValueError("alias_offset_bytes must be a compile-time integer")
+        if alias_offset_bytes < 0:
+            raise ValueError("alias_offset_bytes must be non-negative")
+
     layout = tl._unwrap_if_constexpr(layout)
     if layout is not None and not isinstance(layout, tle.shared_layout):
         # Handle constexpr None
@@ -356,7 +375,11 @@ def alloc(
             layout_handle = layout.to_ir(_semantic.builder)
 
         if storage == tle.smem:
-            if init_value is not None:
+            if alias is not None:
+                alias_ty = _semantic.builder.get_memdesc_type(full_shape, elem_type, layout_handle, "smem")
+                tensor_handle = _semantic.builder.create_memdesc_alias(
+                    alias_ty, alias.handle, alias_offset_bytes)
+            elif init_value is not None:
                 mutable_ty = _semantic.builder.get_memdesc_type(full_shape, elem_type, layout_handle, "smem")
                 tensor_handle = _semantic.builder.create_local_alloc(mutable_ty, init_value.handle)
             else:

test_alias_e2e.py

@@ -0,0 +1,61 @@
+"""E2E test: verify alloc alias shares the same physical smem via modify-and-observe."""
+import torch
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+BLOCK = 128
+
+
+@triton.jit
+def alias_e2e_kernel(in_ptr, out_ptr, N, BLOCK: tl.constexpr):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < N
+
+    # 阶段 1: 分配 v_smem，写入第一批数据 (原始输入)
+    v_smem = tle.gpu.alloc([BLOCK], dtype=tl.float32, scope=tle.gpu.smem)
+    tle.gpu.copy(in_ptr + offs, v_smem, [BLOCK])
+
+    # 阶段 2: alias — o_smem 复用 v_smem 的物理内存
+    o_smem = tle.gpu.alloc(
+        [BLOCK], dtype=tl.float32, scope=tle.gpu.smem,
+        alias=v_smem, alias_offset_bytes=0,
+    )
+
+    # 阶段 3: 关键 — 通过 v_smem 写入第二批数据 (input + BLOCK 偏移)，
+    #         覆盖同一块物理内存。如果 alias 正确，o_smem 应该看到覆盖后的值。
+    offs2 = offs + BLOCK
+    mask2 = offs2 < N
+    tle.gpu.copy(in_ptr + offs2, v_smem, [BLOCK])
+
+    # 阶段 4: 从 o_smem 读出
+    tle.gpu.copy(o_smem, out_ptr + offs, [BLOCK])
+
+
+if __name__ == "__main__":
+    N = 2048  # enough for two BLOCK-sized chunks
+    x = torch.randn(N, device="cuda", dtype=torch.float32)
+    y = torch.zeros(1024, device="cuda", dtype=torch.float32)  # only first 1024
+
+    grid = (triton.cdiv(1024, BLOCK),)
+    alias_e2e_kernel[grid](x, y, 1024, BLOCK=BLOCK)  # 注意：N=1024，第二批数据在 idx 128..1151
+    torch.cuda.synchronize()
+
+    # 验证: y 应该等于 x[128:1152] (第二批)，而不是 x[0:1024] (第一批)
+    # 这证明 v_smem 的覆盖写入被 o_smem 观察到了
+    expected = x[128:1152]
+    max_diff_alias = (expected - y).abs().max().item()
+
+    # 对照组: y 不应该等于第一批数据
+    first_batch = x[:1024]
+    max_diff_original = (first_batch - y).abs().max().item()
+
+    if max_diff_alias < 1e-5 and max_diff_original > 1e-5:
+        print("PASSED: alias E2E verified")
+        print(f"  - o_smem matches overwritten data (second batch): max diff = {max_diff_alias:.2e}")
+        print(f"  - o_smem differs from original data (first batch):  max diff = {max_diff_original:.2e}")
+    else:
+        print(f"FAILED:")
+        print(f"  diff vs overwritten (should be 0): {max_diff_alias:.2e}")
+        print(f"  diff vs original (should be >0):   {max_diff_original:.2e}")

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp

@@ -184,6 +184,8 @@ struct ConvertTritonGPUToLLVM
           typeConverter, patterns, targetInfo, benefit);
       mlir::triton::tle::populateMemDescWGMMAViewOpToLLVMPatterns(
           typeConverter, patterns, benefit);
+      mlir::triton::tle::populateMemDescAliasOpToLLVMPatterns(
+          typeConverter, patterns, benefit);
       mlir::triton::tle::populateExclusiveCumsumOpToLLVMPatterns(
           typeConverter, targetInfo, patterns, benefit);
       mlir::triton::tle::populateWGMMASharedOperandFenceOpToLLVMPatterns(

third_party/tle/dialect/include/IR/TleOps.td

@@ -116,6 +116,31 @@ def Tle_MemDescWGMMAViewOp : Tle_Op<"memdesc_wgmma_view", [Pure,
   let hasVerifier = 1;
 }
 
+def Tle_MemDescAliasOp : Tle_Op<"memdesc_alias", [Pure, MemDescViewTrait]> {
+  let summary = "create a typed view that aliases an existing shared memory descriptor";
+
+  let description = [{
+    This operation returns a descriptor-only view into an existing shared memory
+    descriptor. It does not allocate or move shared memory.
+
+    The result type carries the logical shape, element type, and encoding used
+    by consumers. The source descriptor provides the backing storage. The
+    offset_bytes attribute is a static byte offset from the source descriptor's
+    current view base.
+  }];
+
+  let arguments = (
+    ins TTG_MemDescType:$src,
+    I64Attr:$offset_bytes
+  );
+
+  let results = (outs TTG_MemDescType:$result);
+
+  let assemblyFormat = "$src attr-dict `:` qualified(type($src)) `->` qualified(type($result))";
+
+  let hasVerifier = 1;
+}
+
 def Tle_PipeCreateOp : Tle_Op<"pipe.create"> {
   let summary = "create a typed TLE GPU pipe";
   let arguments = (ins