diff --git a/third_party/ascend/lib/TritonToLinalg/LoadStoreConverter.cpp b/third_party/ascend/lib/TritonToLinalg/LoadStoreConverter.cpp index fe692c81ac..87ac054bb0 100644 --- a/third_party/ascend/lib/TritonToLinalg/LoadStoreConverter.cpp +++ b/third_party/ascend/lib/TritonToLinalg/LoadStoreConverter.cpp @@ -76,6 +76,56 @@ using namespace triton; const std::string MayImplicitTransposeWithLastAxisTAG = "MayImplicitTransposeWithLastAxis"; +namespace { + +Value getRemappedOrOriginal(Value value, ConversionPatternRewriter &rewriter) { + if (Value remapped = rewriter.getRemappedValue(value)) + return remapped; + return value; +} + +bool hasStaticZeroStride(triton::MakeTensorPtrOp makeTensorPtrOp) { + return llvm::any_of(makeTensorPtrOp.getStrides(), [](Value stride) { + auto constantStride = getConstantIntValue(stride); + return constantStride.has_value() && constantStride.value() == 0; + }); +} + +SmallVector getBoundarySizesFromMakeTensorPtr( + triton::MakeTensorPtrOp makeTensorPtrOp, + llvm::ArrayRef boundaryCheck, llvm::ArrayRef tileShape, + const Location &loc, ConversionPatternRewriter &rewriter) { + assert(makeTensorPtrOp.getShape().size() == tileShape.size()); + assert(makeTensorPtrOp.getOffsets().size() == tileShape.size()); + + SmallVector boundarySizes = + getAsIndexOpFoldResult(rewriter.getContext(), tileShape); + const OpFoldResult zero = rewriter.getIndexAttr(0); + + for (size_t i = 0; i < tileShape.size(); ++i) { + if (llvm::find(boundaryCheck, i) == boundaryCheck.end()) + continue; + + OpFoldResult shape = getOpFoldResultOfLayoutInfo( + getRemappedOrOriginal(makeTensorPtrOp.getShape()[i], rewriter), + rewriter); + OpFoldResult offset = getOpFoldResultOfLayoutInfo( + getRemappedOrOriginal(makeTensorPtrOp.getOffsets()[i], rewriter), + rewriter); + OpFoldResult nonNegativeOffset = + maxOpFoldResult(offset, zero, loc, rewriter); + OpFoldResult remaining = maxOpFoldResult( + subOpFoldResult(shape, nonNegativeOffset, loc, rewriter), zero, loc, + rewriter); + boundarySizes[i] = + minOpFoldResult(boundarySizes[i], remaining, loc, rewriter); + } + + return boundarySizes; +} + +} // namespace + LogicalResult AddPtrConverter::matchAndRewrite(triton::AddPtrOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { @@ -325,13 +375,19 @@ LogicalResult LoadConverter::replaceMaskedLoadWithTensorOther( LogicalResult LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { + auto makeTensorPtrOp = op.getPtr().getDefiningOp(); + const bool hasZeroStrideMTP = + makeTensorPtrOp && hasStaticZeroStride(makeTensorPtrOp); // Check if tt.load is modified by AddPtrConverter to a specified state. - if (checkModifiedByAddPtrConverter(op).succeeded()) { + if (!hasZeroStrideMTP && checkModifiedByAddPtrConverter(op).succeeded()) { return continueModifyFromAddPtrConverter(op, adaptor, rewriter); } - auto ptr = adaptor.getPtr(); + Value ptr = hasZeroStrideMTP ? rewriter.getRemappedValue(op.getPtr()) + : adaptor.getPtr(); + if (!ptr) + return rewriter.notifyMatchFailure(op, "missing remapped tensor pointer"); auto mask = op.getMask(); auto other = op.getOther(); auto loc = op.getLoc(); @@ -367,7 +423,7 @@ LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor, } int64_t lastStride = -1; - if (isa(ptr)) { + if (!hasZeroStrideMTP && isa(ptr)) { auto u = ptr; while (auto blkArg = dyn_cast(u)) { if (auto forOp = dyn_cast(blkArg.getOwner()->getParentOp())) { @@ -392,14 +448,15 @@ LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor, return rewriter.notifyMatchFailure( op, "LoadOp expects a memref, not a memref of pointers"); } - if (!op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) { + if (!hasZeroStrideMTP && + !op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) { auto memrefOp = dyn_cast(ptr.getDefiningOp()); auto ret = mlir::ConverterUtils::getLastStrideOfReinterpretCastOp(memrefOp); if (ret.has_value()) lastStride = *ret; } bool mayImplicitTransposeWithLastAxis = - (existDotFlag) && + (!hasZeroStrideMTP) && (existDotFlag) && (!op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) && (lastStride != 1 && mlir::ConverterUtils::isaPermutedMemRefType(memRefType)); @@ -464,9 +521,12 @@ LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor, // boundary check auto boundaryCheck = op.getBoundaryCheck(); if (!boundaryCheck.empty()) { - auto makeTensorPtrOp = op.getPtr().getDefiningOp(); - auto boundarySizes = mlir::ConverterUtils::getBoundarySizes( - boundaryCheck, /*remapped*/ ptr, loc, rewriter); + auto boundarySizes = + hasZeroStrideMTP + ? getBoundarySizesFromMakeTensorPtr(makeTensorPtrOp, boundaryCheck, + memRefShape, loc, rewriter) + : mlir::ConverterUtils::getBoundarySizes( + boundaryCheck, /*remapped*/ ptr, loc, rewriter); // handle the padding auto padding = op.getPadding(); SmallVector srcOffsets(boundarySizes.size(), diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/zero_stride_block_ptr_load.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/zero_stride_block_ptr_load.mlir new file mode 100644 index 0000000000..bdf82dd736 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/zero_stride_block_ptr_load.mlir @@ -0,0 +1,128 @@ +// RUN: triton-opt "--triton-to-linalg=global-kernel=false named-ops=True" --split-input-file %s | FileCheck %s + +// A direct MTP load with zero strides must derive its boundary from the MTP +// logical shape/offset. The converted physical offset is always zero and +// cannot be divided by a zero physical stride to recover that information. +// CHECK-LABEL: func.func @zero_stride_dynamic +// CHECK: memref.reinterpret_cast {{.*}}strides: [0, 0] +// CHECK: arith.subi +// CHECK: arith.maxsi +// CHECK: arith.minsi +// CHECK: memref.copy +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @zero_stride_dynamic( + %src: !tt.ptr {tt.divisibility = 16 : i32}, + %dst: !tt.ptr {tt.divisibility = 16 : i32}, + %shape_m: i32, %shape_n: i32, %offset_m: i32, %offset_n: i32) { + %c0_i32 = arith.constant 0 : i32 + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c4_i64 = arith.constant 4 : i64 + %shape_m_i64 = arith.extsi %shape_m : i32 to i64 + %shape_n_i64 = arith.extsi %shape_n : i32 to i64 + %src_block = tt.make_tensor_ptr %src, [%shape_m_i64, %shape_n_i64], + [%c0_i64, %c0_i64], [%offset_m, %offset_n] + {order = array} : > + %value = tt.load %src_block {boundaryCheck = array, padding = 1 : i32} + : !tt.ptr> + %dst_block = tt.make_tensor_ptr %dst, [%c4_i64, %c4_i64], + [%c4_i64, %c1_i64], [%c0_i32, %c0_i32] + {order = array} : > + tt.store %dst_block, %value : !tt.ptr> + tt.return + } +} + +// ----- + +// CHECK-LABEL: func.func @mixed_zero_stride_dynamic +// CHECK: memref.reinterpret_cast {{.*}}strides: [0, 1] +// CHECK: arith.subi +// CHECK: arith.maxsi +// CHECK: arith.minsi +// CHECK: memref.copy +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @mixed_zero_stride_dynamic( + %src: !tt.ptr {tt.divisibility = 16 : i32}, + %dst: !tt.ptr {tt.divisibility = 16 : i32}, + %shape_m: i32, %shape_n: i32, %offset_m: i32, %offset_n: i32) { + %c0_i32 = arith.constant 0 : i32 + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c4_i64 = arith.constant 4 : i64 + %shape_m_i64 = arith.extsi %shape_m : i32 to i64 + %shape_n_i64 = arith.extsi %shape_n : i32 to i64 + %src_block = tt.make_tensor_ptr %src, [%shape_m_i64, %shape_n_i64], + [%c0_i64, %c1_i64], [%offset_m, %offset_n] + {order = array} : > + %value = tt.load %src_block {boundaryCheck = array, padding = 1 : i32} + : !tt.ptr> + %dst_block = tt.make_tensor_ptr %dst, [%c4_i64, %c4_i64], + [%c4_i64, %c1_i64], [%c0_i32, %c0_i32] + {order = array} : > + tt.store %dst_block, %value : !tt.ptr> + tt.return + } +} + +// ----- + +// CHECK-LABEL: func.func @zero_stride_negative_offset +// CHECK: memref.reinterpret_cast {{.*}}strides: [0, 0] +// CHECK: memref.subview {{.*}}[2, 0] +// CHECK: memref.copy +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @zero_stride_negative_offset( + %src: !tt.ptr {tt.divisibility = 16 : i32}, + %dst: !tt.ptr {tt.divisibility = 16 : i32}) { + %cneg2_i32 = arith.constant -2 : i32 + %c0_i32 = arith.constant 0 : i32 + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c4_i64 = arith.constant 4 : i64 + %c5_i64 = arith.constant 5 : i64 + %c6_i64 = arith.constant 6 : i64 + %src_block = tt.make_tensor_ptr %src, [%c6_i64, %c5_i64], + [%c0_i64, %c0_i64], [%cneg2_i32, %c0_i32] + {order = array} : > + %value = tt.load %src_block {boundaryCheck = array, padding = 1 : i32} + : !tt.ptr> + %dst_block = tt.make_tensor_ptr %dst, [%c4_i64, %c4_i64], + [%c4_i64, %c1_i64], [%c0_i32, %c0_i32] + {order = array} : > + tt.store %dst_block, %value : !tt.ptr> + tt.return + } +} + +// ----- + +// A direct MTP with statically nonzero strides must keep the legacy physical +// offset reconstruction path. +// CHECK-LABEL: func.func @nonzero_stride_dynamic +// CHECK: arith.divsi +// CHECK: arith.remsi +// CHECK: memref.copy +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @nonzero_stride_dynamic( + %src: !tt.ptr {tt.divisibility = 16 : i32}, + %dst: !tt.ptr {tt.divisibility = 16 : i32}, + %shape_m: i32, %shape_n: i32, %offset_m: i32, %offset_n: i32) { + %c0_i32 = arith.constant 0 : i32 + %c1_i64 = arith.constant 1 : i64 + %c4_i64 = arith.constant 4 : i64 + %c8_i64 = arith.constant 8 : i64 + %shape_m_i64 = arith.extsi %shape_m : i32 to i64 + %shape_n_i64 = arith.extsi %shape_n : i32 to i64 + %src_block = tt.make_tensor_ptr %src, [%shape_m_i64, %shape_n_i64], + [%c8_i64, %c1_i64], [%offset_m, %offset_n] + {order = array} : > + %value = tt.load %src_block {boundaryCheck = array, padding = 1 : i32} + : !tt.ptr> + %dst_block = tt.make_tensor_ptr %dst, [%c4_i64, %c4_i64], + [%c4_i64, %c1_i64], [%c0_i32, %c0_i32] + {order = array} : > + tt.store %dst_block, %value : !tt.ptr> + tt.return + } +}