[ssbuffer](feat) extend new scenario in splitmatmul

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<SplitInfo> 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<linalg::MatmulOp>(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<linalg::MatmulOp>(hivm::traceDefOp<linalg::MatmulOp>(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<SplitInfo> 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<SplitInfo> 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<SplitInfo> 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<scf::ForOp>(splitInfo.outerOutValue.getDefiningOp());
+    if (!forOp) {
+        return splitInfo.outerOutValue;
+    }
+
+    auto outputType = dyn_cast<RankedTensorType>(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<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, ub, lb);
+    Value zeroValue;
+    if (auto floatType = dyn_cast<FloatType>(elmType)) {
+        APFloat zeroAPFloat = APFloat::getZero(floatType.getFloatSemantics());
+        zeroValue = rewriter.create<arith::ConstantFloatOp>(loc, zeroAPFloat, floatType).getResult();
+    } else if (auto intType = dyn_cast<IntegerType>(elmType)) {
+        zeroValue = rewriter.create<arith::ConstantIntOp>(loc, 0, intType).getResult();
+    }
+    auto fillOp = rewriter.create<linalg::FillOp>(loc, zeroValue, splitInfo.outerOutValue);
+    auto selectOp = rewriter.create<arith::SelectOp>(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();

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<?xi32>, %[[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<?xi32> 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<?xi32>, %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<?xi32> 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>
+  }
 }