[TLE][MTHREADS] Fix shared memory alias conflict for tle.gpu.alloc buffers

third_party/mthreads/lib/Analysis/Alias.cpp

@@ -3,6 +3,10 @@
 #include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Support/LLVM.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#ifdef __TLE__
+#include "triton/Dialect/Triton/IR/Types.h"
+#include "llvm/ADT/STLExtras.h"
+#endif
 
 namespace mlir {
 
@@ -19,9 +23,32 @@ AliasInfo AliasInfo::join(const AliasInfo &lhs, const AliasInfo &rhs) {
   return ret;
 }
 
+#ifdef __TLE__
+static bool isTritonPtrLikeType(Type type) {
+  if (isa<triton::PointerType>(type))
+    return true;
+  if (auto tensorTy = dyn_cast<RankedTensorType>(type))
+    return isa<triton::PointerType>(tensorTy.getElementType());
+  return false;
+}
+
+static AliasInfo
+joinOperandAliases(ArrayRef<const dataflow::Lattice<AliasInfo> *> operands) {
+  AliasInfo aliasInfo;
+  for (auto *operand : operands)
+    aliasInfo = AliasInfo::join(aliasInfo, operand->getValue());
+  return aliasInfo;
+}
+#endif // __TLE__
+
 LogicalResult SharedMemoryAliasAnalysis::visitOperation(
     Operation *op, ArrayRef<const dataflow::Lattice<AliasInfo> *> operands,
     ArrayRef<dataflow::Lattice<AliasInfo> *> results) {
+#ifdef __TLE__
+  if (results.empty())
+    return success();
+#endif // __TLE__
+
   AliasInfo aliasInfo;
   bool pessimistic = true;
   auto result = op->getResult(0);
@@ -39,6 +66,13 @@ LogicalResult SharedMemoryAliasAnalysis::visitOperation(
   } else if (op->hasTrait<OpTrait::MemDescViewTrait>()) {
     aliasInfo = AliasInfo(operands[0]->getValue());
     pessimistic = false;
+#ifdef __TLE__
+  } else if (op->getName().getStringRef() == "musa_tle.local_pointers" &&
+             !operands.empty()) {
+    // Treat local pointer views as aliases of their source memdesc.
+    aliasInfo = operands[0]->getValue();
+    pessimistic = false;
+#endif // __TLE__
   } else if (isa<ub::PoisonOp>(op)) {
     aliasInfo = AliasInfo();
     pessimistic = false;
@@ -48,7 +82,25 @@ LogicalResult SharedMemoryAliasAnalysis::visitOperation(
   }
 
   if (pessimistic) {
+#ifdef __TLE__
+    // Propagate aliases through pointer-producing ops such as
+    // tt.splat / tt.broadcast / tt.addptr so that shared buffers
+    // allocated via tle.gpu.alloc stay live across pointer arithmetic
+    // users (tl.load / tl.store through derived pointers).
+    bool propagated = false;
+    for (auto [idx, result] : llvm::enumerate(results)) {
+      Value value = op->getResult(idx);
+      if (isTritonPtrLikeType(value.getType())) {
+        AliasInfo ptrAlias = joinOperandAliases(operands);
+        propagateIfChanged(result, result->join(ptrAlias));
+        propagated = true;
+      }
+    }
+    if (!propagated)
+      setAllToEntryStates(results);
+#else
     setAllToEntryStates(results);
+#endif
     return success();
   }
   // Join all lattice elements

third_party/mthreads/lib/Analysis/Membar.cpp

@@ -34,8 +34,17 @@ bool shouldTrackMusaSquadDotOp(Operation *op) {
 AllocationSlice::AllocationSlice(Value value,
                                  Interval<size_t> allocationInterval)
     : allocationInterval(allocationInterval) {
+#ifdef __TLE__
+  // With TLE, alias propagation may associate pointer-type values (e.g.
+  // derived from tle.local_pointers / tt.addptr) with shared memory buffers.
+  // Use dyn_cast so we degrade gracefully instead of asserting.
+  this->accessTy = dyn_cast<triton::gpu::MemDescType>(value.getType());
+  if (!accessTy)
+    return;
+#else
   auto accessTy = cast<triton::gpu::MemDescType>(value.getType());
   this->accessTy = accessTy;
+#endif
 
   // Get the memdesc_subslice information if present. If no subslice is
   // present the whole interval is accessed
@@ -364,9 +373,16 @@ void MembarAnalysis::update(Operation *op, BlockInfo *blockInfo,
 
     if (!curBlockInfo.syncReadSlices.empty() ||
         !curBlockInfo.syncWriteSlices.empty()) {
+#ifdef __TLE__
+      // Some scratch-buffer ops can also carry explicit shared-memory
+      // effects (e.g. tle.gpu.local_ptr accesses). Keep conservative
+      // dependency tracking instead of hard-failing here; the normal
+      // barrier-insertion logic below will handle overlaps.
+#else
       llvm::report_fatal_error(
           "scratch buffer operations should not have any shared memory "
           "dependencies");
+#endif
     }
     auto interval = allocation->getAllocatedInterval(scratchBufferId);
     auto scratchSlice = AllocationSlice(interval);

third_party/mthreads/python/test/unit/tle/test_tle_local_ptr.py

@@ -827,3 +827,255 @@ def test_tle_local_ptr_atomic_cas_runtime_round_trip():
     _local_ptr_atomic_cas_update_kernel[(1, )](out, num_warps=1)
 
     torch.testing.assert_close(out.cpu(), torch.tensor([9], dtype=torch.int32), rtol=0, atol=0)
+
+
+# ---------------------------------------------------------------------------
+# Shared-memory alias / lifetime regression tests
+#
+# These tests verify that a tle.gpu.alloc buffer accessed through
+# tle.gpu.local_ptr + pointer arithmetic is NOT overwritten by scratch
+# shared-memory allocations (e.g. layout conversions from tl.trans or
+# tl.device_print) that occur between two loads from the same buffer.
+#
+# The scratch-triggering operation is placed *between* the two loads so
+# that the shared-memory allocator must keep the histogram buffer live
+# across the scratch allocation.  Without proper alias propagation the
+# allocator reuses the histogram's smem region for the scratch buffer,
+# corrupting the second load.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _local_ptr_smem_alias_trans_zero_init_kernel(
+    logits_ptr,
+    out_c1_ptr,
+    out_c2_ptr,
+    out_trans_ptr,
+    stride0,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_BINS: tl.constexpr,
+    VEC: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    logits_ptr += row_id * stride0
+
+    s_histogram = tle.gpu.alloc(
+        (NUM_BINS, ),
+        dtype=tl.int32,
+        nv_mma_shared_layout=False,
+    )
+    s_histogram_ptr = tle.gpu.local_ptr(s_histogram, (0, ))
+
+    tl.debug_barrier()
+    tl.store(s_histogram_ptr + tl.arange(0, NUM_BINS), 0)
+    tl.debug_barrier()
+
+    lane = tl.arange(0, BLOCK_SIZE)
+    vec = tl.arange(0, VEC)
+    base = BLOCK_SIZE * VEC + lane * VEC
+    offs = base[:, None] + vec[None, :]
+    x_vec = tl.load(logits_ptr + offs)
+
+    tl.debug_barrier()
+    my_counts1 = tl.load(s_histogram_ptr + tl.arange(0, NUM_BINS))
+    tl.debug_barrier()
+
+    # tl.trans + store placed BETWEEN the two histogram loads.
+    # The store forces the convert_layout (scratch shared memory) to
+    # execute here, while s_histogram is still live.  Without proper
+    # alias propagation the scratch overlaps the histogram region and
+    # corrupts the second load below.
+    transposed = tl.trans(x_vec)
+    trans_offs = vec[:, None] * BLOCK_SIZE + lane[None, :]
+    tl.store(out_trans_ptr + trans_offs, transposed)
+    tl.debug_barrier()
+
+    my_counts2 = tl.load(s_histogram_ptr + tl.arange(0, NUM_BINS))
+    tl.debug_barrier()
+
+    tl.store(out_c1_ptr + tl.arange(0, NUM_BINS), my_counts1)
+    tl.store(out_c2_ptr + tl.arange(0, NUM_BINS), my_counts2)
+
+
+@triton.jit
+def _local_ptr_smem_alias_trans_pattern_init_kernel(
+    logits_ptr,
+    out_c1_ptr,
+    out_c2_ptr,
+    out_trans_ptr,
+    stride0,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_BINS: tl.constexpr,
+    VEC: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    logits_ptr += row_id * stride0
+
+    s_histogram = tle.gpu.alloc(
+        (NUM_BINS, ),
+        dtype=tl.int32,
+        nv_mma_shared_layout=False,
+    )
+    s_histogram_ptr = tle.gpu.local_ptr(s_histogram, (0, ))
+
+    pattern = tl.arange(0, NUM_BINS).to(tl.int32) * 7 + 3
+
+    tl.debug_barrier()
+    tl.store(s_histogram_ptr + tl.arange(0, NUM_BINS), pattern)
+    tl.debug_barrier()
+
+    lane = tl.arange(0, BLOCK_SIZE)
+    vec = tl.arange(0, VEC)
+    base = BLOCK_SIZE * VEC + lane * VEC
+    offs = base[:, None] + vec[None, :]
+    x_vec = tl.load(logits_ptr + offs)
+
+    tl.debug_barrier()
+    my_counts1 = tl.load(s_histogram_ptr + tl.arange(0, NUM_BINS))
+    tl.debug_barrier()
+
+    transposed = tl.trans(x_vec)
+    trans_offs = vec[:, None] * BLOCK_SIZE + lane[None, :]
+    tl.store(out_trans_ptr + trans_offs, transposed)
+    tl.debug_barrier()
+
+    my_counts2 = tl.load(s_histogram_ptr + tl.arange(0, NUM_BINS))
+    tl.debug_barrier()
+
+    tl.store(out_c1_ptr + tl.arange(0, NUM_BINS), my_counts1)
+    tl.store(out_c2_ptr + tl.arange(0, NUM_BINS), my_counts2)
+
+
+@triton.jit
+def _local_ptr_smem_alias_trans_between_loads_control_kernel(
+    logits_ptr,
+    out_c1_ptr,
+    out_c2_ptr,
+    out_trans_ptr,
+    stride0,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_BINS: tl.constexpr,
+    VEC: tl.constexpr,
+):
+    """Control: same trans+store between loads but histogram accessed via
+    global-memory pointer (no tle.gpu.alloc), so alias propagation is not
+    needed and the test should always pass."""
+    row_id = tl.program_id(0)
+    logits_ptr += row_id * stride0
+
+    lane = tl.arange(0, BLOCK_SIZE)
+    vec = tl.arange(0, VEC)
+    base = BLOCK_SIZE * VEC + lane * VEC
+    offs = base[:, None] + vec[None, :]
+    x_vec = tl.load(logits_ptr + offs)
+
+    hist_offs = tl.arange(0, NUM_BINS)
+    my_counts1 = tl.load(out_c1_ptr + hist_offs)
+    tl.debug_barrier()
+
+    transposed = tl.trans(x_vec)
+    trans_offs = vec[:, None] * BLOCK_SIZE + lane[None, :]
+    tl.store(out_trans_ptr + trans_offs, transposed)
+    tl.debug_barrier()
+
+    my_counts2 = tl.load(out_c1_ptr + hist_offs)
+    tl.debug_barrier()
+
+    tl.store(out_c2_ptr + hist_offs, my_counts2)
+
+
+@pytest.mark.skipif(not torch.musa.is_available(), reason="MUSA device is not available")
+def test_tle_local_ptr_smem_alias_not_overwritten_by_trans_zero_init():
+    """Regression: a tle.gpu.alloc buffer initialised to zero must stay
+    zero when a tl.trans convert_layout (scratch shared memory) occurs
+    between two loads from the same buffer."""
+    torch.manual_seed(789)
+    num_bins = 2048
+    block_size = 512
+    vec = 4
+    vocab_size = 129280
+
+    logits = torch.randn(1, vocab_size, device="musa", dtype=torch.float32)
+    c1 = torch.full((num_bins, ), -1, device="musa", dtype=torch.int32)
+    c2 = torch.full((num_bins, ), -1, device="musa", dtype=torch.int32)
+    out_trans = torch.empty((1, vocab_size), device="musa", dtype=torch.float32)
+
+    _local_ptr_smem_alias_trans_zero_init_kernel[(1, )](
+        logits,
+        c1,
+        c2,
+        out_trans,
+        logits.stride(0),
+        BLOCK_SIZE=block_size,
+        NUM_BINS=num_bins,
+        VEC=vec,
+        num_warps=16,
+    )
+
+    expected = torch.zeros(num_bins, dtype=torch.int32)
+    torch.testing.assert_close(c1.cpu(), expected, rtol=0, atol=0)
+    torch.testing.assert_close(c2.cpu(), expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not torch.musa.is_available(), reason="MUSA device is not available")
+def test_tle_local_ptr_smem_alias_not_overwritten_by_trans_pattern_init():
+    """Same as above but with a known non-zero pattern to make corruption
+    obvious."""
+    torch.manual_seed(789)
+    num_bins = 2048
+    block_size = 512
+    vec = 4
+    vocab_size = 129280
+
+    logits = torch.randn(1, vocab_size, device="musa", dtype=torch.float32)
+    c1 = torch.full((num_bins, ), -1, device="musa", dtype=torch.int32)
+    c2 = torch.full((num_bins, ), -1, device="musa", dtype=torch.int32)
+    out_trans = torch.empty((1, vocab_size), device="musa", dtype=torch.float32)
+
+    _local_ptr_smem_alias_trans_pattern_init_kernel[(1, )](
+        logits,
+        c1,
+        c2,
+        out_trans,
+        logits.stride(0),
+        BLOCK_SIZE=block_size,
+        NUM_BINS=num_bins,
+        VEC=vec,
+        num_warps=16,
+    )
+
+    expected = torch.arange(num_bins, dtype=torch.int32) * 7 + 3
+    torch.testing.assert_close(c1.cpu(), expected, rtol=0, atol=0)
+    torch.testing.assert_close(c2.cpu(), expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not torch.musa.is_available(), reason="MUSA device is not available")
+def test_tle_local_ptr_smem_alias_trans_between_loads_control():
+    """Control: same trans+store between loads but without tle.gpu.alloc
+    (global-memory buffer instead of shared memory).  Verifies the trans
+    itself doesn't corrupt data when there's no smem-alias issue."""
+    torch.manual_seed(789)
+    num_bins = 2048
+    block_size = 512
+    vec = 4
+    vocab_size = 129280
+
+    logits = torch.randn(1, vocab_size, device="musa", dtype=torch.float32)
+    c1 = torch.zeros(num_bins, device="musa", dtype=torch.int32)
+    c2 = torch.full((num_bins, ), -1, device="musa", dtype=torch.int32)
+    out_trans = torch.empty((1, vocab_size), device="musa", dtype=torch.float32)
+
+    _local_ptr_smem_alias_trans_between_loads_control_kernel[(1, )](
+        logits,
+        c1,
+        c2,
+        out_trans,
+        logits.stride(0),
+        BLOCK_SIZE=block_size,
+        NUM_BINS=num_bins,
+        VEC=vec,
+        num_warps=16,
+    )
+
+    expected = torch.zeros(num_bins, dtype=torch.int32)
+    torch.testing.assert_close(c2.cpu(), expected, rtol=0, atol=0)