diff --git a/third_party/mthreads/lib/Analysis/Alias.cpp b/third_party/mthreads/lib/Analysis/Alias.cpp index e2d0c90499..eceb0f74bc 100644 --- a/third_party/mthreads/lib/Analysis/Alias.cpp +++ b/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(type)) + return true; + if (auto tensorTy = dyn_cast(type)) + return isa(tensorTy.getElementType()); + return false; +} + +static AliasInfo +joinOperandAliases(ArrayRef *> operands) { + AliasInfo aliasInfo; + for (auto *operand : operands) + aliasInfo = AliasInfo::join(aliasInfo, operand->getValue()); + return aliasInfo; +} +#endif // __TLE__ + LogicalResult SharedMemoryAliasAnalysis::visitOperation( Operation *op, ArrayRef *> operands, ArrayRef *> 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()) { 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(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 diff --git a/third_party/mthreads/lib/Analysis/Membar.cpp b/third_party/mthreads/lib/Analysis/Membar.cpp index 5fe72f3da3..d2c1cb5c8d 100644 --- a/third_party/mthreads/lib/Analysis/Membar.cpp +++ b/third_party/mthreads/lib/Analysis/Membar.cpp @@ -34,8 +34,17 @@ bool shouldTrackMusaSquadDotOp(Operation *op) { AllocationSlice::AllocationSlice(Value value, Interval 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(value.getType()); + if (!accessTy) + return; +#else auto accessTy = cast(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); diff --git a/third_party/mthreads/python/test/unit/tle/test_tle_local_ptr.py b/third_party/mthreads/python/test/unit/tle/test_tle_local_ptr.py index 8ac21e576a..ff565ec83d 100644 --- a/third_party/mthreads/python/test/unit/tle/test_tle_local_ptr.py +++ b/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)