From 1dc67914133ae5dc67a364c7feb53a53383fd57f Mon Sep 17 00:00:00 2001 From: fishofnanqi <1074959344@qq.com> Date: Wed, 5 Aug 2026 12:42:08 +0000 Subject: [PATCH] [ssbuffer](feat) extend new scenario in splitmatmul Signed-off-by: fishofnanqi <1074959344@qq.com> --- .../StandardizeOp/SplitMatmulPattern.cpp | 90 +++++++++++++++++-- .../StandardizeOp/split_matmul.mlir | 52 +++++++++++ 2 files changed, 137 insertions(+), 5 deletions(-) diff --git a/third_party/ascend/lib/DynamicCVPipeline/StandardizeOp/SplitMatmulPattern.cpp b/third_party/ascend/lib/DynamicCVPipeline/StandardizeOp/SplitMatmulPattern.cpp index cf47770198..6d7c8852f1 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/StandardizeOp/SplitMatmulPattern.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/StandardizeOp/SplitMatmulPattern.cpp @@ -83,6 +83,7 @@ struct SplitInfo { Value outerInValue; Value outerOutValue; bool shouldSplit; + bool supported = false; }; } // namespace @@ -522,6 +523,31 @@ static std::optional handleMayNotExec(linalg::MatmulOp matmulOp) } auto initVal = forOp.getTiedLoopInit(blockArg)->get(); auto result = forOp.getTiedLoopResult(blockArg); + + // Check if matmul is used by subsequent L0C + auto matchMatmulC = [](Operation *op, Value value) { + if (auto nextMatmulOp = dyn_cast(op)) { + auto inputs = parseMatmulInputs(nextMatmulOp); + return inputs.a != value && inputs.b != value && inputs.bias == value; + } + return false; + }; + auto usedByL0C = traceChainUser(result, true, matchMatmulC, [](Operation *op, Value value) { return false; }); + + // Check if used by L0C in different block + if (usedByL0C.has_value() && usedByL0C.value()->getBlock() != result.getParentBlock()) { + return SplitInfo {false, initVal, result, false, false}; + } + + // Check if outerInValue is defined by another matmul in different block + auto defMatmul = dyn_cast_if_present(hivm::traceDefOp(initVal).value_or(nullptr)); + if (defMatmul) { + auto defInMatmulBlock = CVPipeline::getAncestorInBlock(defMatmul, matmulOp->getBlock()); + if (!defInMatmulBlock) { + return SplitInfo {true, initVal, result, false, true}; + } + } + return SplitInfo {true, initVal, result, true}; } @@ -552,7 +578,7 @@ static std::optional shouldSplit(linalg::MatmulOp matmulOp, bool need auto matmulInput = parseMatmulInputs(matmulOp); if (needSplitAll) { LOG_DEBUG("Split because needSplitAll is true. " << matmulOp); - return SplitInfo {false, matmulInput.bias, matmulOp.getResult(0), true}; + return SplitInfo {false, matmulInput.bias, matmulOp.getResult(0), true, false}; } bool argsLimitedInMatmul = true; @@ -564,7 +590,7 @@ static std::optional shouldSplit(linalg::MatmulOp matmulOp, bool need } if (!argsLimitedInMatmul) { LOG_DEBUG("Split because bias is not limited in args" << matmulOp); // S25 - return SplitInfo {mayNotExec, outerInValue, outerOutValue, true}; + return SplitInfo {mayNotExec, outerInValue, outerOutValue, true, false}; } if (mayNotExec) { @@ -577,10 +603,60 @@ static std::optional shouldSplit(linalg::MatmulOp matmulOp, bool need if (!shouldSplitByInput(matmulOp, outerOutValue, outerInValue) && !shouldSplitByOutput(matmulOp, outerOutValue, outerInValue)) { - return SplitInfo {mayNotExec, outerInValue, outerOutValue, false}; + return SplitInfo {mayNotExec, outerInValue, outerOutValue, false, false}; } - return SplitInfo {mayNotExec, outerInValue, outerOutValue, true}; + return SplitInfo {mayNotExec, outerInValue, outerOutValue, true, false}; +} + +/** + * Handles the mayNotExec case by creating select operations. + * This function creates a select operation to choose between the actual result + * and a zero-filled result based on whether the loop will execute. + * + * Returns the select result if mayNotExec is handled, otherwise returns the original outerOutValue. + */ +static Value handleMayNotExecSelect(linalg::MatmulOp matmulOp, PatternRewriter &rewriter, + SplitInfo &splitInfo, bool replaceUses = true) +{ + auto forOp = llvm::dyn_cast_if_present(splitInfo.outerOutValue.getDefiningOp()); + if (!forOp) { + return splitInfo.outerOutValue; + } + + auto outputType = dyn_cast(parseMatmulInputs(matmulOp).bias.getType()); + if (!outputType) { + return splitInfo.outerOutValue; + } + auto elmType = outputType.getElementType(); + Location loc = matmulOp.getLoc(); + + rewriter.setInsertionPointAfterValue(splitInfo.outerOutValue); + + auto lb = forOp.getLowerBound(); + auto ub = forOp.getUpperBound(); + + Value executed = rewriter.create(loc, arith::CmpIPredicate::sgt, ub, lb); + Value zeroValue; + if (auto floatType = dyn_cast(elmType)) { + APFloat zeroAPFloat = APFloat::getZero(floatType.getFloatSemantics()); + zeroValue = rewriter.create(loc, zeroAPFloat, floatType).getResult(); + } else if (auto intType = dyn_cast(elmType)) { + zeroValue = rewriter.create(loc, 0, intType).getResult(); + } + auto fillOp = rewriter.create(loc, zeroValue, splitInfo.outerOutValue); + auto selectOp = rewriter.create(loc, executed, splitInfo.outerOutValue, fillOp.getResult(0)); + + if (replaceUses) { + // Replace uses of outerOutValue with select result, except for the select and fill operations + splitInfo.outerOutValue.replaceUsesWithIf( + selectOp.getResult(), + [&](OpOperand &operand) { return operand.getOwner() != selectOp && operand.getOwner() != fillOp; }); + } + + forOp->setAttr(CVPipeline::kHIVMMatmulLimitedInCubeAttr, rewriter.getUnitAttr()); + + return selectOp.getResult(); } static LogicalResult splitMatmul(linalg::MatmulOp matmulOp, PatternRewriter &rewriter, SplitInfo splitInfo) @@ -716,7 +792,11 @@ LogicalResult SplitMatmulPattern::matchAndRewrite(linalg::MatmulOp matmulOp, Pat } if (splitInfo.mayNotExec) { - matmulOp->setAttr(CVPipeline::kMayNotExec, rewriter.getUnitAttr()); + if (splitInfo.supported) { + handleMayNotExecSelect(matmulOp, rewriter, splitInfo); + } else { + matmulOp->setAttr(CVPipeline::kMayNotExec, rewriter.getUnitAttr()); + } } return success(); diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/StandardizeOp/split_matmul.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/StandardizeOp/split_matmul.mlir index 340fb8baaa..a6c56bc946 100644 --- a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/StandardizeOp/split_matmul.mlir +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/StandardizeOp/split_matmul.mlir @@ -446,5 +446,57 @@ module { } return %result : tensor<32x32xf32> } + + // Case 14: There are for-loop cascading and the upper bound is loaded from gm. + // While upper bound is loaded from gm, the for-loop may not execute, + // it need to add SelectOp when the result of for-loop is used; + // However, non-final for loops within cascaded for loops cannot be wrapped with selectOp, + // otherwise, they will not be identified to remain in L0C. + + // CHECK-LABEL: func.func @case14_cascaded_for_loops_maynotexec + // CHECK-SAME: (%[[GM:.*]]: memref, %[[A:.*]]: tensor<32x64xf32>, %[[B:.*]]: tensor<64x32xf32>) -> tensor<32x32xf32> + // CHECK-DAG: %[[OFFSET:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[LB:.*]] = arith.constant 0 : i32 + // CHECK-DAG: %[[CST_ONE:.*]] = arith.constant 1 : i32 + // CHECK-DAG: %[[CST_ZERO:.*]] = arith.constant 0.000000e+00 : f32 + + // CHECK: %[[REINTERPRET:.*]] = memref.reinterpret_cast %[[GM]] to offset: [%[[OFFSET]]], sizes: [1], strides: [1] : memref to memref<1xi32, strided<[1], offset: ?>> + // CHECK: %[[UB:.*]] = memref.load %[[REINTERPRET]][%[[OFFSET]]] : memref<1xi32, strided<[1], offset: ?>> + + // CHECK: %[[FOR_RESULT1:.*]] = scf.for %{{.*}} = %{{.*}} to %[[UB]] step %{{.*}} iter_args(%[[ITER_BIAS1:.*]] = %{{.*}}) -> (tensor<32x32xf32>) : i32 { + // CHECK: %[[MM1:.*]] = linalg.matmul {ssbuffer.loop_carried_l0c} ins(%[[A]], %[[B]] : tensor<32x64xf32>, tensor<64x32xf32>) outs(%[[ITER_BIAS1]] : tensor<32x32xf32>) -> tensor<32x32xf32> + // CHECK: scf.yield %[[MM1]] : tensor<32x32xf32> + // CHECK: } + // CHECK: %[[FOR_RESULT2:.*]] = scf.for %{{.*}} = %{{.*}} to %[[UB]] step %{{.*}} iter_args(%[[ITER_BIAS2:.*]] = %[[FOR_RESULT1]]) -> (tensor<32x32xf32>) : i32 { + // CHECK: %[[MM2:.*]] = linalg.matmul {ssbuffer.loop_carried_l0c} ins(%[[A]], %[[B]] : tensor<32x64xf32>, tensor<64x32xf32>) outs(%[[ITER_BIAS2]] : tensor<32x32xf32>) -> tensor<32x32xf32> + // CHECK: scf.yield %[[MM2]] : tensor<32x32xf32> + // CHECK: } {hivm.matmul_limited_in_cube} + + // CHECK: %[[COND:.*]] = arith.cmpi sgt, %[[UB]], %[[LB]] : i32 + // CHECK: %[[ZERO_TENSOR:.*]] = linalg.fill ins(%{{.*}} : f32) outs(%[[FOR_RESULT2]] : tensor<32x32xf32>) -> tensor<32x32xf32> + // CHECK: %[[RESULT:.*]] = arith.select %[[COND]], %[[FOR_RESULT2]], %[[ZERO_TENSOR]] : tensor<32x32xf32> + // CHECK: return %[[RESULT]] : tensor<32x32xf32> + func.func @case14_cascaded_for_loops_maynotexec(%gm: memref, %A: tensor<32x64xf32>, %B: tensor<64x32xf32>) -> tensor<32x32xf32> { + %c0_idx = arith.constant 0 : index + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + + %cst_zero = arith.constant 0.0 : f32 + %empty1 = tensor.empty() : tensor<32x32xf32> + %c = linalg.fill ins(%cst_zero : f32) outs(%empty1 : tensor<32x32xf32>) -> tensor<32x32xf32> + %reinterpret_cast = memref.reinterpret_cast %gm to offset: [%c0_idx], sizes: [1], strides: [1] : memref to memref<1xi32, strided<[1], offset: ?>> + %upper_bound = memref.load %reinterpret_cast[%c0_idx] : memref<1xi32, strided<[1], offset: ?>> + + // Chain 1: zero initialization, each matmul in a for loop + %for1_result = scf.for %i1 = %c0_i32 to %upper_bound step %c1_i32 iter_args(%bias1 = %c) -> (tensor<32x32xf32>) : i32 { + %mm1 = linalg.matmul ins(%A, %B : tensor<32x64xf32>, tensor<64x32xf32>) outs(%bias1 : tensor<32x32xf32>) -> tensor<32x32xf32> + scf.yield %mm1 : tensor<32x32xf32> + } + %for2_result = scf.for %i2 = %c0_i32 to %upper_bound step %c1_i32 iter_args(%bias2 = %for1_result) -> (tensor<32x32xf32>) : i32 { + %mm2 = linalg.matmul ins(%A, %B : tensor<32x64xf32>, tensor<64x32xf32>) outs(%bias2 : tensor<32x32xf32>) -> tensor<32x32xf32> + scf.yield %mm2 : tensor<32x32xf32> + } + func.return %for2_result : tensor<32x32xf32> + } }