diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp index 0a336b9644..8ba85883fe 100644 --- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp +++ b/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()) { + value = alias.getSrc(); + continue; + } if (auto trans = value.getDefiningOp()) { value = trans.getSrc(); continue; diff --git a/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp b/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp index 4a1e023e5a..a3b751d8f7 100644 --- a/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp +++ b/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp @@ -84,6 +84,10 @@ static Value stripMemDescViews(Value value) { value = subslice.getSrc(); continue; } + if (auto alias = value.getDefiningOp()) { + value = alias.getSrc(); + continue; + } if (auto trans = value.getDefiningOp()) { value = trans.getSrc(); continue; diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.cpp index b3e14a8a9f..226ff6d718 100644 --- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.cpp +++ b/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()) { + value = alias.getSrc(); + continue; + } if (auto trans = value.getDefiningOp()) { value = trans.getSrc(); continue; @@ -548,6 +552,10 @@ static MemDescResource getMemDescResource(Value value) { value = subslice.getSrc(); continue; } + if (auto alias = value.getDefiningOp()) { + value = alias.getSrc(); + continue; + } if (auto trans = value.getDefiningOp()) { value = trans.getSrc(); continue; diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp index 11cbfa73c2..f3a0961c7c 100644 --- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp +++ b/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 dotCanBeProperlyAsync(ttng::WarpGroupDotOp dotOp, // allowed in between. Value transitiveOperand = operand; while (isa_and_nonnull( + ttg::MemDescReshapeOp, ttg::MemDescSubsliceOp, + mlir::triton::tle::MemDescAliasOp>( transitiveOperand.getDefiningOp()) || isa(transitiveOperand)) { auto blockArg = dyn_cast(transitiveOperand); diff --git a/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp b/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp index c736697046..4666563817 100644 --- a/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp +++ b/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" diff --git a/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/lib/Dialect/TritonGPU/Transforms/Utility.cpp index aae234cb86..c4c7ed5d41 100644 --- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp +++ b/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::MemDescReshapeOp, ttg::MemDescSubsliceOp, + tle::MemDescAliasOp>( transitiveOperand.getDefiningOp()) || isa(transitiveOperand)) { if (auto blockArg = dyn_cast(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(user)) { + ttg::MemDescType oldType = alias.getType(); + bool isMutable = cast(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(user)) { newVal = ttg::MemDescTransOp::create(builder, trans.getLoc(), val, trans.getOrder()); diff --git a/python/test/tle/unit/test_tle_alloc_alias.py b/python/test/tle/unit/test_tle_alloc_alias.py new file mode 100644 index 0000000000..bcdb77aaa9 --- /dev/null +++ b/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, + ) diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py index 1680ae3094..8b3d0516bc 100644 --- a/python/triton/experimental/tle/language/gpu/core.py +++ b/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: diff --git a/test_alias_e2e.py b/test_alias_e2e.py new file mode 100644 index 0000000000..ba33c8328d --- /dev/null +++ b/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}") diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index 552dc79edf..0848d67636 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/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( diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td index a0a7b233eb..fd0601923c 100644 --- a/third_party/tle/dialect/include/IR/TleOps.td +++ b/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 diff --git a/third_party/tle/dialect/include/Transforms/PatternTleToLLVM.h b/third_party/tle/dialect/include/Transforms/PatternTleToLLVM.h index d759272ee2..4a5d0f4301 100644 --- a/third_party/tle/dialect/include/Transforms/PatternTleToLLVM.h +++ b/third_party/tle/dialect/include/Transforms/PatternTleToLLVM.h @@ -44,6 +44,10 @@ void populateMemDescWGMMAViewOpToLLVMPatterns( mlir::LLVMTypeConverter &typeConverter, mlir::RewritePatternSet &patterns, unsigned benefit = 1); +void populateMemDescAliasOpToLLVMPatterns( + mlir::LLVMTypeConverter &typeConverter, mlir::RewritePatternSet &patterns, + unsigned benefit = 1); + void populateWGMMASharedOperandFenceOpToLLVMPatterns( mlir::LLVMTypeConverter &typeConverter, mlir::RewritePatternSet &patterns, unsigned benefit = 1); diff --git a/third_party/tle/dialect/lib/Analysis/TleMemoryEffectAnalysis.cpp b/third_party/tle/dialect/lib/Analysis/TleMemoryEffectAnalysis.cpp index 666ba0f391..b0efe3e80d 100644 --- a/third_party/tle/dialect/lib/Analysis/TleMemoryEffectAnalysis.cpp +++ b/third_party/tle/dialect/lib/Analysis/TleMemoryEffectAnalysis.cpp @@ -79,6 +79,10 @@ static Value getSharedMemDescRoot(Value value) { current = stripConvertLayouts(subslice.getSrc()); continue; } + if (auto alias = current.getDefiningOp()) { + current = stripConvertLayouts(alias.getSrc()); + continue; + } if (auto reinterpret = current.getDefiningOp()) { current = stripConvertLayouts(reinterpret.getSrc()); continue; diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/ViewOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/ViewOpToLLVM.cpp index 8c8026388b..7748e9ce32 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/ViewOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/ViewOpToLLVM.cpp @@ -56,6 +56,38 @@ struct MemDescWGMMAViewOpConversion } }; +struct MemDescAliasOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + + LogicalResult + matchAndRewrite(triton::tle::MemDescAliasOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + auto b = TritonLLVMOpBuilder(loc, rewriter); + auto srcTy = op.getSrc().getType(); + auto resultTy = op.getType(); + auto srcElemTy = getTypeConverter()->convertType(srcTy.getElementType()); + auto resultElemTy = + getTypeConverter()->convertType(resultTy.getElementType()); + + auto srcSmemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), + srcElemTy, rewriter); + Value base = srcSmemObj.getShmemAffineBase(loc, rewriter, srcTy); + int64_t offsetBytes = op.getOffsetBytesAttr().getInt(); + if (offsetBytes != 0) + base = b.gep(base.getType(), i8_ty, base, b.i32_val(offsetBytes)); + + auto dstSmemObj = + SharedMemoryObject(base, resultElemTy, resultTy.getRank(), loc, + rewriter); + auto retVal = + LLVM::getStructFromSharedMemoryObject(loc, dstSmemObj, rewriter); + rewriter.replaceOp(op, retVal); + return success(); + } +}; + } // namespace void mlir::triton::tle::populateMemDescWGMMAViewOpToLLVMPatterns( @@ -63,3 +95,9 @@ void mlir::triton::tle::populateMemDescWGMMAViewOpToLLVMPatterns( unsigned benefit) { patterns.add(typeConverter, benefit); } + +void mlir::triton::tle::populateMemDescAliasOpToLLVMPatterns( + LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, + unsigned benefit) { + patterns.add(typeConverter, benefit); +} diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp index 3349b6446f..05ca014728 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -186,6 +186,61 @@ LogicalResult MemDescWGMMAViewOp::verify() { return success(); } +static std::optional +getStaticMemDescByteSize(triton::gpu::MemDescType type) { + int64_t numElements = 1; + for (int64_t dim : type.getShape()) { + if (ShapedType::isDynamic(dim) || dim < 0) + return std::nullopt; + if (dim != 0 && numElements > std::numeric_limits::max() / dim) + return std::nullopt; + numElements *= dim; + } + + int64_t elementBits = type.getElementTypeBitWidth(); + int64_t elementBytes = (elementBits + 7) / 8; + if (elementBytes <= 0) + return std::nullopt; + if (numElements > std::numeric_limits::max() / elementBytes) + return std::nullopt; + return numElements * elementBytes; +} + +LogicalResult MemDescAliasOp::verify() { + auto srcType = getSrc().getType(); + auto resultType = getType(); + int64_t offsetBytes = getOffsetBytesAttr().getInt(); + + if (srcType.getMemorySpace() != resultType.getMemorySpace()) + return emitOpError("expects source and result memory spaces to match"); + if (!isa(srcType.getMemorySpace())) + return emitOpError("expects shared memory descriptors"); + if (resultType.getMutableMemory() && !srcType.getMutableMemory()) + return emitOpError("cannot create a mutable alias from an immutable source"); + if (offsetBytes < 0) + return emitOpError("expects non-negative offset_bytes"); + if (offsetBytes > std::numeric_limits::max()) + return emitOpError("expects offset_bytes to fit in i32 for shared memory " + "lowering"); + + int64_t resultElementBytes = (resultType.getElementTypeBitWidth() + 7) / 8; + if (resultElementBytes <= 0) + return emitOpError("expects byte-addressable result element type"); + if (offsetBytes % resultElementBytes != 0) + return emitOpError("expects offset_bytes to be aligned to the result " + "element byte width"); + + std::optional srcBytes = getStaticMemDescByteSize(srcType); + std::optional resultBytes = getStaticMemDescByteSize(resultType); + if (!srcBytes || !resultBytes) + return emitOpError("expects static source and result memdesc byte sizes"); + if (*resultBytes > *srcBytes || + offsetBytes > *srcBytes - *resultBytes) + return emitOpError("result byte range must fit within the source view"); + + return success(); +} + LogicalResult WGMMASharedOperandFenceOp::verify() { if (getDeps().empty()) return emitOpError("expects at least one shared-memory operand"); diff --git a/third_party/tle/dialect/lib/Transforms/EncodingRematerialization.cpp b/third_party/tle/dialect/lib/Transforms/EncodingRematerialization.cpp index 451e446626..67144a195c 100644 --- a/third_party/tle/dialect/lib/Transforms/EncodingRematerialization.cpp +++ b/third_party/tle/dialect/lib/Transforms/EncodingRematerialization.cpp @@ -24,6 +24,7 @@ #include "tle/dialect/include/Transforms/EncodingRematerialization.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "tle/dialect/include/IR/Dialect.h" #include "mlir/IR/IRMapping.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "triton/Dialect/Triton/IR/Dialect.h" @@ -117,6 +118,10 @@ static Value getSharedMemDescRoot(Value value) { current = stripConvertLayouts(subslice.getSrc()); continue; } + if (auto alias = current.getDefiningOp()) { + current = stripConvertLayouts(alias.getSrc()); + continue; + } if (auto reinterpret = current.getDefiningOp()) { current = stripConvertLayouts(reinterpret.getSrc()); continue; diff --git a/third_party/tle/dialect/lib/Transforms/TleLowerPipeToNvws.cpp b/third_party/tle/dialect/lib/Transforms/TleLowerPipeToNvws.cpp index 3897109f89..5034f11762 100644 --- a/third_party/tle/dialect/lib/Transforms/TleLowerPipeToNvws.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleLowerPipeToNvws.cpp @@ -248,6 +248,10 @@ static Value getMemDescRoot(Value value) { current = canonicalizePipeField(subslice.getSrc()); continue; } + if (auto alias = current.getDefiningOp()) { + current = canonicalizePipeField(alias.getSrc()); + continue; + } break; } return current; @@ -554,6 +558,10 @@ getCommitFieldRootForStore(Value memdesc, PipeWriterCommitOp commit) { current = canonicalizePipeField(subslice.getSrc()); continue; } + if (auto alias = current.getDefiningOp()) { + current = canonicalizePipeField(alias.getSrc()); + continue; + } break; } diff --git a/third_party/tle/dialect/lib/Transforms/TleMaterializeTileStylePipeline.cpp b/third_party/tle/dialect/lib/Transforms/TleMaterializeTileStylePipeline.cpp index a0895ce6dd..c2131465c5 100644 --- a/third_party/tle/dialect/lib/Transforms/TleMaterializeTileStylePipeline.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleMaterializeTileStylePipeline.cpp @@ -167,6 +167,10 @@ static Value stripProducerMemDescViews(Value value) { current = subslice.getSrc(); continue; } + if (auto alias = current.getDefiningOp()) { + current = alias.getSrc(); + continue; + } if (auto index = current.getDefiningOp()) { current = index.getSrc(); continue; @@ -177,7 +181,7 @@ static Value stripProducerMemDescViews(Value value) { } static bool isTileProducerViewLikeOp(Operation *op) { - return isa(op) || + return isa(op) || op->getName().getStringRef() == "tle.memdesc_wgmma_view"; } @@ -215,6 +219,18 @@ static Operation *cloneWithUpdatedMemDescViewType(OpBuilder &builder, newOp->setAttrs(op->getAttrs()); return mapResults(newOp); } + if (auto alias = dyn_cast(op)) { + Value src = mapping.lookupOrDefault(alias.getSrc()); + auto oldTy = alias.getType(); + bool isMutable = cast(src.getType()).getMutableMemory(); + auto newTy = ttg::MemDescType::get( + oldTy.getShape(), oldTy.getElementType(), oldTy.getEncoding(), + oldTy.getMemorySpace(), isMutable, oldTy.getAllocShape()); + auto newOp = MemDescAliasOp::create( + builder, alias.getLoc(), newTy, src, alias.getOffsetBytesAttr()); + newOp->setAttrs(op->getAttrs()); + return mapResults(newOp); + } if (auto trans = dyn_cast(op)) { Value src = mapping.lookupOrDefault(trans.getSrc()); auto newOp = ttg::MemDescTransOp::create(builder, trans.getLoc(), src, diff --git a/third_party/tle/dialect/lib/Transforms/TleOptimizeLocalPointerAsyncStores.cpp b/third_party/tle/dialect/lib/Transforms/TleOptimizeLocalPointerAsyncStores.cpp index 72991f83da..0c2450a95a 100644 --- a/third_party/tle/dialect/lib/Transforms/TleOptimizeLocalPointerAsyncStores.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleOptimizeLocalPointerAsyncStores.cpp @@ -211,6 +211,10 @@ static Value getMemDescRoot(Value value) { current = subslice.getSrc(); continue; } + if (auto alias = current.getDefiningOp()) { + current = alias.getSrc(); + continue; + } break; } return current; diff --git a/third_party/tle/dialect/lib/Transforms/TleScheduleTmaStoreSync.cpp b/third_party/tle/dialect/lib/Transforms/TleScheduleTmaStoreSync.cpp index 8a3804a8b9..3c0555df12 100644 --- a/third_party/tle/dialect/lib/Transforms/TleScheduleTmaStoreSync.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleScheduleTmaStoreSync.cpp @@ -81,6 +81,10 @@ static Value getMemDescRoot(Value value) { current = canonicalizeWarpSpecializeCapture(subslice.getSrc()); continue; } + if (auto alias = current.getDefiningOp()) { + current = canonicalizeWarpSpecializeCapture(alias.getSrc()); + continue; + } if (auto trans = current.getDefiningOp()) { current = canonicalizeWarpSpecializeCapture(trans.getSrc()); continue; diff --git a/third_party/tle/dialect/lib/Transforms/TleTileStylePipelineSchedule.cpp b/third_party/tle/dialect/lib/Transforms/TleTileStylePipelineSchedule.cpp index 3679f459c6..db952b1ede 100644 --- a/third_party/tle/dialect/lib/Transforms/TleTileStylePipelineSchedule.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleTileStylePipelineSchedule.cpp @@ -92,6 +92,10 @@ static Value stripProducerMemDescViews(Value value) { current = subslice.getSrc(); continue; } + if (auto alias = current.getDefiningOp()) { + current = alias.getSrc(); + continue; + } if (auto index = current.getDefiningOp()) { current = index.getSrc(); continue; @@ -110,7 +114,7 @@ static bool isTleOp(Operation *op) { } static bool isTileProducerViewLikeOp(Operation *op) { - return isa(op) || + return isa(op) || op->getName().getStringRef() == "tle.memdesc_wgmma_view"; } diff --git a/third_party/tle/test/GPU/test_tle_memdesc_alias.mlir b/third_party/tle/test/GPU/test_tle_memdesc_alias.mlir new file mode 100644 index 0000000000..a7c003527f --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_memdesc_alias.mlir @@ -0,0 +1,69 @@ +// RUN: triton-opt %s -split-input-file --verify-diagnostics | FileCheck %s + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @valid_memdesc_alias + tt.func @valid_memdesc_alias(%arg0: tensor<64x64xbf16, #blocked>) { + %src = ttg.local_alloc %arg0 : (tensor<64x64xbf16, #blocked>) -> !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> + // CHECK: tle.memdesc_alias + // CHECK-SAME: offset_bytes = 128 : i64 + %alias = tle.memdesc_alias %src {offset_bytes = 128 : i64} : !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> -> !ttg.memdesc<32x64xbf16, #shared1, #smem, mutable> + tt.return + } +} + +// ----- + +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + tt.func @reject_negative_offset(%src: !ttg.memdesc<64x64xbf16, #shared, #smem, mutable>) { + // expected-error @+1 {{expects non-negative offset_bytes}} + %alias = tle.memdesc_alias %src {offset_bytes = -1 : i64} : !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> -> !ttg.memdesc<1x1xbf16, #shared, #smem, mutable> + tt.return + } +} + +// ----- + +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + tt.func @reject_unaligned_offset(%src: !ttg.memdesc<64x64xbf16, #shared, #smem, mutable>) { + // expected-error @+1 {{expects offset_bytes to be aligned}} + %alias = tle.memdesc_alias %src {offset_bytes = 1 : i64} : !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> -> !ttg.memdesc<1x1xbf16, #shared, #smem, mutable> + tt.return + } +} + +// ----- + +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + tt.func @reject_out_of_bounds(%src: !ttg.memdesc<64x64xbf16, #shared, #smem, mutable>) { + // expected-error @+1 {{result byte range must fit within the source view}} + %alias = tle.memdesc_alias %src {offset_bytes = 8192 : i64} : !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> -> !ttg.memdesc<1x1xbf16, #shared, #smem, mutable> + tt.return + } +} + +// ----- + +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + tt.func @reject_mutable_alias_from_immutable(%src: !ttg.memdesc<64x64xbf16, #shared, #smem>) { + // expected-error @+1 {{cannot create a mutable alias from an immutable source}} + %alias = tle.memdesc_alias %src {offset_bytes = 0 : i64} : !ttg.memdesc<64x64xbf16, #shared, #smem> -> !ttg.memdesc<1x1xbf16, #shared, #smem, mutable> + tt.return + } +} diff --git a/third_party/tle/test/GPU/test_tle_memdesc_alias_allocation.mlir b/third_party/tle/test/GPU/test_tle_memdesc_alias_allocation.mlir new file mode 100644 index 0000000000..c954c094cc --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_memdesc_alias_allocation.mlir @@ -0,0 +1,19 @@ +// RUN: triton-opt %s --allocate-shared-memory-nv='compute-capability=90 ptx-version=81' | FileCheck %s + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + // CHECK: ttg.shared = 8192 : i32 + tt.func @alias_does_not_allocate_new_smem(%arg0: tensor<64x64xbf16, #blocked>) { + // CHECK: ttg.local_alloc + // CHECK-SAME: allocation.offset = 0 : i32 + %src = ttg.local_alloc %arg0 : (tensor<64x64xbf16, #blocked>) -> !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> + // CHECK: tle.memdesc_alias + // CHECK-NOT: allocation.offset = 8192 : i32 + %alias = tle.memdesc_alias %src {offset_bytes = 0 : i64} : !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> -> !ttg.memdesc<32x64xbf16, #shared1, #smem, mutable> + tt.return + } +} diff --git a/third_party/tle/test/GPU/test_tle_memdesc_alias_to_llvm.mlir b/third_party/tle/test/GPU/test_tle_memdesc_alias_to_llvm.mlir new file mode 100644 index 0000000000..7c32dfe454 --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_memdesc_alias_to_llvm.mlir @@ -0,0 +1,18 @@ +// RUN: triton-opt %s --allocate-shared-memory-nv='compute-capability=90 ptx-version=81' --convert-triton-gpu-to-llvm='compute-capability=90 ptx-version=81' -reconcile-unrealized-casts | FileCheck %s + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: llvm.func @memdesc_alias_to_llvm + tt.func @memdesc_alias_to_llvm(%arg0: tensor<64x64xbf16, #blocked>) { + %src = ttg.local_alloc %arg0 : (tensor<64x64xbf16, #blocked>) -> !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> + %alias = tle.memdesc_alias %src {offset_bytes = 128 : i64} : !ttg.memdesc<64x64xbf16, #shared, #smem, mutable> -> !ttg.memdesc<1x1xbf16, #shared, #smem, mutable> + %value = ttg.local_load %alias : !ttg.memdesc<1x1xbf16, #shared, #smem, mutable> -> tensor<1x1xbf16, #blocked> + tt.return + } +} + +// CHECK: llvm.getelementptr +// CHECK-NOT: tle.memdesc_alias diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 29e99e1d12..6d313b1b9f 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -370,6 +370,13 @@ void init_triton_tle_ir(py::module &&m) { return self.create(resultType, src, offsets); }) + .def("create_memdesc_alias", + [](TritonOpBuilder &self, Type resultType, Value src, + int64_t offsetBytes) -> Value { + return self.create( + resultType, src, + self.getBuilder().getI64IntegerAttr(offsetBytes)); + }) .def("create_warp_return", [](TritonOpBuilder &self) -> Operation * { return self.create();