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/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp b/third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp index 483b468a26..740443d196 100644 --- a/third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp +++ b/third_party/ascend/lib/TritonToLinalg/StridedLoadStoreRewrite.cpp @@ -81,6 +81,13 @@ static std::optional getStaticConstInt(Value v) { return std::nullopt; } +static bool hasStaticZeroStride(ValueRange strides) { + return llvm::any_of(strides, [](Value stride) { + auto value = getStaticConstInt(stride); + return value.has_value() && value.value() == 0; + }); +} + static std::optional getStaticMaskUpperBound(Value mask) { if (!mask) return std::nullopt; @@ -817,40 +824,43 @@ static LogicalResult tryRewriteBlockPtrLoad(triton::LoadOp op, ArrayRef shape = resultType.getShape(); int64_t rank = static_cast(shape.size()); - // ---- order must match the "non-permuted" layout: order[i] == rank-1-i, - // i.e. innermost (fastest-changing) is the last dim of the tensor. - // ImplicitPermute handles permuted layouts via tt.trans, so we - // leave anything non-canonical alone. - auto order = mtpt.getOrder(); - if (static_cast(order.size()) != rank) - return failure(); - for (int64_t i = 0; i < rank; ++i) { - if (order[i] != rank - 1 - i) { - return failure(); - } - } - // ---- stride check ---- auto strides = mtpt.getStrides(); if (strides.empty() || static_cast(strides.size()) != rank) return failure(); - // Stride dispatch: strided DMA on the MTE engine only supports power-of-two - // strides; a non-power-of-two stride would degrade to a slow scalar access. - // Dynamic strides stay on the structured SIMD path because they may be - // runtime stride 1 or power-of-two, where SIMT stride is slower. So we - // only rewrite to SIMT stride for *static non-power-of-two* strides: - // stride 1 -> contiguous; stride 2 (even dim) -> deinterleave; - // stride >= 4 (power of two) -> (compact) strided DMA. - APInt lastStrideC; - if (!matchPattern(strides.back(), m_ConstantInt(&lastStrideC))) - return failure(); - int64_t lastStride = std::abs(lastStrideC.getSExtValue()); - if (lastStride <= 1) - return failure(); - if (lastStride == 2) - return failure(); // even -> deinterleave; odd -> strided DMA - if ((lastStride & (lastStride - 1)) == 0) - return failure(); // power-of-two >= 4 -> strided DMA + const bool zeroStrideBroadcast = hasStaticZeroStride(strides); + int64_t lastStride = 0; + if (!zeroStrideBroadcast) { + // ---- order must match the "non-permuted" layout: order[i] == rank-1-i, + // i.e. innermost (fastest-changing) is the last dim of the tensor. + // ImplicitPermute handles permuted layouts via tt.trans, so we + // leave anything non-canonical alone. + auto order = mtpt.getOrder(); + if (static_cast(order.size()) != rank) + return failure(); + for (int64_t i = 0; i < rank; ++i) { + if (order[i] != rank - 1 - i) + return failure(); + } + + // Stride dispatch: strided DMA on the MTE engine only supports power-of-two + // strides; a non-power-of-two stride would degrade to a slow scalar access. + // Dynamic strides stay on the structured SIMD path because they may be + // runtime stride 1 or power-of-two, where SIMT stride is slower. So we + // only rewrite to SIMT stride for *static non-power-of-two* strides: + // stride 1 -> contiguous; stride 2 (even dim) -> deinterleave; + // stride >= 4 (power of two) -> (compact) strided DMA. + APInt lastStrideC; + if (!matchPattern(strides.back(), m_ConstantInt(&lastStrideC))) + return failure(); + lastStride = std::abs(lastStrideC.getSExtValue()); + if (lastStride <= 1) + return failure(); + if (lastStride == 2) + return failure(); // even -> deinterleave; odd -> strided DMA + if ((lastStride & (lastStride - 1)) == 0) + return failure(); // power-of-two >= 4 -> strided DMA + } // ---- Compute per-axis effective base offsets: mtpt.offsets[d] + // (advance.offsets[d] if present) @@ -915,8 +925,8 @@ static LogicalResult tryRewriteBlockPtrLoad(triton::LoadOp op, scalarBase = rewriter.create(loc, scalarBase, prod); } - if (resultType.getRank() >= 1 && resultType.getRank() <= 3 && - resultType.hasStaticShape()) { + if (!zeroStrideBroadcast && resultType.getRank() >= 1 && + resultType.getRank() <= 3 && resultType.hasStaticShape()) { FailureOr> numels = getBlockPtrStrideNumels( op.getOperation(), op.getMask(), op.getBoundaryCheck(), mtpt, resultType, logicalOffsets, rewriter); @@ -944,7 +954,7 @@ static LogicalResult tryRewriteBlockPtrLoad(triton::LoadOp op, } } - if (rank <= 3) + if (!zeroStrideBroadcast && rank <= 3) return failure(); auto i64TensorTy = RankedTensorType::get(shape, i64Ty); diff --git a/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp b/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp index eacd50eddc..92f08aaa96 100644 --- a/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp +++ b/third_party/ascend/lib/TritonToUnstructure/UnstructureConversionPass.cpp @@ -75,6 +75,45 @@ static Type getResultElementType(RankedTensorType ptrType) { return resultElementType; } +static triton::MakeTensorPtrOp getBaseMakeTensorPtr(Value ptr) { + while (auto advance = ptr.getDefiningOp()) + ptr = advance.getPtr(); + return ptr.getDefiningOp(); +} + +static bool hasStaticZeroStride(triton::MakeTensorPtrOp makeTensorPtr) { + return makeTensorPtr && + llvm::any_of(makeTensorPtr.getStrides(), [](Value stride) { + auto constantStride = getConstantIntValue(stride); + return constantStride.has_value() && constantStride.value() == 0; + }); +} + +static Value castToIndex(Value value, Location loc, PatternRewriter &rewriter) { + if (value.getType().isIndex()) + return value; + return rewriter.create(loc, rewriter.getIndexType(), + value); +} + +static Value getPaddingValue(triton::LoadOp op, Type elementType, + PatternRewriter &rewriter) { + auto padding = op.getPadding(); + if (!padding.has_value()) + return Value(); + + auto loc = op.getLoc(); + TypedAttr padAttr = rewriter.getZeroAttr(elementType); + if (padding.value() == triton::PaddingOption::PAD_NAN) { + auto floatType = dyn_cast(elementType); + if (!floatType) + return Value(); + auto nan = llvm::APFloat::getNaN(floatType.getFloatSemantics()); + padAttr = rewriter.getFloatAttr(elementType, nan); + } + return rewriter.create(loc, padAttr); +} + static int64_t getTypeSizeInByte(Type type) { if (auto intType = dyn_cast(type)) { return intType.getWidth() / kBitsPerByte; @@ -479,6 +518,25 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( auto ptrOffsetInfo = offsetMap.at(ptr); + // BishengIR does not accept a memref with a static zero stride. A + // make_tensor_ptr load nevertheless has well-defined broadcast semantics: + // the physical offset does not advance on that axis, while boundary_check + // remains in the logical shape/offset coordinate system. Force only this + // load through the existing scalar-loop fallback so it never materializes a + // zero-strided memref. Other accesses keep their original structured path. + triton::MakeTensorPtrOp zeroStrideMakeTensorPtr; + if constexpr (std::is_same_v) { + zeroStrideMakeTensorPtr = getBaseMakeTensorPtr(ptr); + // A5 pure-SIMT has a dedicated earlier rewrite to tt.indirect_load. Keep + // that efficient route intact; every other target/mode takes this generic + // fallback. + if (hasStaticZeroStride(zeroStrideMakeTensorPtr) && + !(compileOn91095Flag && forceSimtTemplateFlag)) + ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank()); + else + zeroStrideMakeTensorPtr = nullptr; + } + if (checkUnstructureAnnotated(op, rewriter)) ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank()); @@ -582,6 +640,14 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( if (isLoadLike) { iterArg = rewriter.create(loc, resultShape, resultElementType); + if constexpr (std::is_same_v) { + if (zeroStrideMakeTensorPtr && !op.getBoundaryCheck().empty()) { + if (Value padding = getPaddingValue(op, resultElementType, rewriter)) { + iterArg = rewriter.create(loc, padding, iterArg) + .getResult(0); + } + } + } } Value newOpResult = nullptr; @@ -606,8 +672,11 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( extractedShape.push_back(size); } else { scf::ForOp forOp; - if (auto mtptOp = - srcPtr.template getDefiningOp()) { + auto mtptOp = + zeroStrideMakeTensorPtr + ? zeroStrideMakeTensorPtr + : srcPtr.template getDefiningOp(); + if (mtptOp) { auto tptShape = mtptOp.getShape()[i]; if (tptShape.getType() != rewriter.getIndexType()) { tptShape = rewriter.create( @@ -618,18 +687,38 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( Value loopLower = zeroIdx; Value loopUpper = sizeVal; + if constexpr (std::is_same_v) { + if (zeroStrideMakeTensorPtr && llvm::find(op.getBoundaryCheck(), i) != + op.getBoundaryCheck().end()) { + Value logicalOffset = + castToIndex(ptrOffsetInfo.getOffsets()[i], loc, rewriter); + Value logicalShape = + castToIndex(zeroStrideMakeTensorPtr.getShape()[i], loc, rewriter); + Value firstValidIndex = + rewriter.create(loc, zeroIdx, logicalOffset); + firstValidIndex = + rewriter.create(loc, firstValidIndex, zeroIdx); + loopLower = + rewriter.create(loc, firstValidIndex, sizeVal); + + Value remaining = + rewriter.create(loc, logicalShape, logicalOffset); + remaining = rewriter.create(loc, remaining, zeroIdx); + loopUpper = rewriter.create(loc, remaining, sizeVal); + } + } if (mstate && i < mstate->dims.size() && i < mstate->offsets.size()) { Value maskOffset = getValueOrCreateConstantIndexOp(rewriter, loc, mstate->offsets[i]); maskOffset = rewriter.create(loc, maskOffset, zeroIdx); maskOffset = rewriter.create(loc, maskOffset, sizeVal); - loopLower = maskOffset; + loopLower = rewriter.create(loc, loopLower, maskOffset); Value maskDim = getValueOrCreateConstantIndexOp(rewriter, loc, mstate->dims[i]); maskDim = rewriter.create(loc, maskOffset, maskDim); maskDim = rewriter.create(loc, maskDim, sizeVal); - loopUpper = maskDim; + loopUpper = rewriter.create(loc, loopUpper, maskDim); } if (isLoadLike) { @@ -658,8 +747,11 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( Value extractedOffset; if (fullyUnstructured) { - if (auto mtptOp = - srcPtr.template getDefiningOp()) { + auto mtptOp = + zeroStrideMakeTensorPtr + ? zeroStrideMakeTensorPtr + : srcPtr.template getDefiningOp(); + if (mtptOp) { auto I64Type = rewriter.getIntegerType(64); srcPtr = mtptOp.getBase(); extractedOffset = rewriter.create(loc, 0, 64); diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/zero_stride_block_ptr_indirect_a5.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/zero_stride_block_ptr_indirect_a5.mlir new file mode 100644 index 0000000000..e9625b5d0f --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/zero_stride_block_ptr_indirect_a5.mlir @@ -0,0 +1,33 @@ +// RUN: triton-opt --pass-pipeline="builtin.module(triton-to-unstructure{compile-on-910-95=true force-simt-template=true},triton-to-linalg{compile-on-910-95=true enable-nd2nz-on-vector=false enable-select-analysis=true global-kernel=false named-ops=true})" %s | FileCheck %s + +// On A5 SIMT, a block-pointer load with a statically zero stride must not +// materialize a zero-strided memref. BishengIR rejects that layout. Route the +// load through per-element offsets instead; the zero stride naturally produces +// a zero offset increment while the logical boundary check remains intact. +// CHECK-LABEL: func.func private @triton_indirect_load +// CHECK-LABEL: func.func @zero_stride_block_ptr_indirect_a5 +// CHECK-NOT: strided<[0 +// CHECK: call @triton_indirect_load +module attributes {hacc.target = #hacc.target<"Ascend910_9589">} { + tt.func public @zero_stride_block_ptr_indirect_a5( + %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 + } +} 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..af712a224f --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/zero_stride_block_ptr_load.mlir @@ -0,0 +1,128 @@ +// RUN: triton-opt --pass-pipeline="builtin.module(triton-to-unstructure{compile-on-910-95=false force-simt-template=true},triton-to-linalg{compile-on-910-95=false enable-nd2nz-on-vector=false enable-select-analysis=true global-kernel=false named-ops=true})" --split-input-file %s | FileCheck %s + +// A block-pointer load with a static zero stride must avoid materializing a +// zero-strided memref. It is lowered through the scalar-loop fallback, which +// loads only the logical in-bounds region and initializes the rest with the +// requested padding. +// CHECK-LABEL: func.func @zero_stride_dynamic +// CHECK: linalg.fill +// CHECK: scf.for +// CHECK: memref.load +// CHECK-NOT: strided<[0 +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: linalg.fill +// CHECK: scf.for +// CHECK: memref.load +// CHECK-NOT: strided<[0 +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: linalg.fill +// CHECK: scf.for +// CHECK: memref.load +// CHECK-NOT: strided<[0 +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 + } +} diff --git a/third_party/ascend/unittest/pytest_ut/test_zero_stride_block_ptr_load.py b/third_party/ascend/unittest/pytest_ut/test_zero_stride_block_ptr_load.py new file mode 100644 index 0000000000..450af8af4c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_zero_stride_block_ptr_load.py @@ -0,0 +1,93 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + + +@triton.jit +def zero_stride_block_ptr_load_kernel( + input_ptr, + output_ptr, + shape_m, + shape_n, + offset_m, + offset_n, + STRIDE_M: tl.constexpr, + STRIDE_N: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + input_block = tl.make_block_ptr( + base=input_ptr, + shape=(shape_m, shape_n), + strides=(STRIDE_M, STRIDE_N), + offsets=(offset_m, offset_n), + block_shape=(BLOCK_M, BLOCK_N), + order=(1, 0), + ) + output_block = tl.make_block_ptr( + base=output_ptr, + shape=(BLOCK_M, BLOCK_N), + strides=(BLOCK_N, 1), + offsets=(0, 0), + block_shape=(BLOCK_M, BLOCK_N), + order=(1, 0), + ) + value = tl.load(input_block, boundary_check=(0, 1), padding_option="zero") + tl.store(output_block, value) + + +def _reference(values, shape, strides, offsets, block_shape): + expected = torch.zeros(block_shape, dtype=values.dtype) + for row in range(block_shape[0]): + for col in range(block_shape[1]): + logical_row = offsets[0] + row + logical_col = offsets[1] + col + if 0 <= logical_row < shape[0] and 0 <= logical_col < shape[1]: + physical_offset = logical_row * strides[0] + logical_col * strides[1] + expected[row, col] = values[physical_offset] + return expected + + +@pytest.mark.parametrize( + "strides,offsets,values", + [ + ((0, 0), (4, 4), torch.tensor([3.5], dtype=torch.float32)), + ((0, 1), (4, 3), torch.arange(5, dtype=torch.float32)), + ((0, 0), (-2, 0), torch.tensor([7.0], dtype=torch.float32)), + ], +) +def test_zero_stride_block_ptr_load(strides, offsets, values): + shape = (6, 5) + block_shape = (4, 4) + input_tensor = values.npu() + output = torch.empty(block_shape, dtype=values.dtype).npu() + + zero_stride_block_ptr_load_kernel[(1, )]( + input_tensor, + output, + shape[0], + shape[1], + offsets[0], + offsets[1], + STRIDE_M=strides[0], + STRIDE_N=strides[1], + BLOCK_M=block_shape[0], + BLOCK_N=block_shape[1], + ) + + expected = _reference(values, shape, strides, offsets, block_shape) + torch.testing.assert_close(output.cpu(), expected)