triton-lang · ohblacks · Sep 1, 2026
@@ -59,16 +59,14 @@
 #include "ascend/include/DynamicCVPipeline/ComputeBlockOpt/Passes.h"
 #include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h"
 
-#include "Common.h"
+#include "ComputeBlockOpt/SplitIfByBlockId/Common.h"
 #include "DynamicCVPipeline/Common/FallbackHelper.h"
 #include "bishengir/Dialect/HIVM/IR/HIVM.h"
 #include "bishengir/Dialect/HIVM/IR/HIVMImpl.h"
 
 static constexpr const char *DEBUG_TYPE = "SplitIfByBlockId";
 static constexpr llvm::StringLiteral kSkippedDeltaformerKernel =
     "parallel_deltaformer_fwd_kernel";
-static constexpr llvm::StringLiteral kSkippedChunkwiseKernel =
-    "chunkwise_fwd_kernel";
 static constexpr llvm::StringLiteral kSkippedParallelNsaFwdKernel =
     "parallel_nsa_fwd_kernel";
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
@@ -162,9 +160,11 @@ struct CandidateIf {
   // value-level cross-group + yield plan
   YieldAugmentation yieldAug;
 
-  bool needsSplit() const {
-    return thenGroups.size() >= 2 || elseGroups.size() >= 2;
-  }
+  // Pre-computed split decisions: ≥2 groups AND mixed matmul/VECTOR
+  bool shouldSplitThen = false;
+  bool shouldSplitElse = false;
+
+  bool needsSplit() const { return shouldSplitThen || shouldSplitElse; }
 };
 
 } // namespace
@@ -183,6 +183,89 @@ static bool hasNestedIfs(const CandidateIf &c) {
   return false;
 }
 
+/// Check whether a BlockGroup contains a linalg.matmul op.
+/// Only linalg::MatmulOp is considered CUBE for the mixed-type filter.
+static bool groupHasMatmul(const BlockGroup &group) {
+  for (auto *op : group.ops) {
+    if (isa<linalg::MatmulOp>(op)) {
+      return true;
+    }
+    auto found =
+        op->walk([](linalg::MatmulOp) { return WalkResult::interrupt(); });
+    if (found.wasInterrupted()) {
+      return true;
+    }
+  }
+  for (auto nestedIf : group.nestedIfs) {
+    auto found = nestedIf->walk(
+        [](linalg::MatmulOp) { return WalkResult::interrupt(); });
+    if (found.wasInterrupted()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+/// Check whether a list of groups forms a splittable CVC/VCV pattern.
+/// Only linalg.matmul defines CUBE for this check.
+/// - Exactly 1 matmul group: must be CVC (VECTOR groups on both sides).
+/// - 2 or more matmul groups (VCV, CVCVC, ...): splittable as long as a
+///   VECTOR group exists.
+static bool hasMixedCoreTypes(const SmallVector<BlockGroup, 0> &groups) {
+  unsigned matmulCount = 0;
+  size_t lastMatmulIdx = 0;
+  for (size_t i = 0; i < groups.size(); ++i) {
+    if (groupHasMatmul(groups[i])) {
+      ++matmulCount;
+      lastMatmulIdx = i;
+    }
+  }
+
+  if (matmulCount == 0) {
+    return false;
+  }
+
+  if (matmulCount == 1) {
+    // CVC: the single CUBE group must have VECTOR groups on both sides.
+    return lastMatmulIdx > 0 && lastMatmulIdx + 1 < groups.size();
+  }
+
+  // matmulCount >= 2 (VCV, CVCVC, ...): split as long as a VECTOR group
+  // exists.
+  for (auto &g : groups) {
+    if (!groupHasMatmul(g)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+/// Merge consecutive groups that share the same core_type
+/// (both CUBE or both VECTOR). This ensures splits follow core_type
+/// transitions: one scf.if per merged segment (3 for CVC/VCV,
+/// 5 for CVCVC, ...).
+static SmallVector<BlockGroup>
+mergeConsecutiveSameCoreType(SmallVector<BlockGroup> groups) {
+  if (groups.size() < 2) {
+    return groups;
+  }
+  SmallVector<BlockGroup> merged;
+  merged.push_back(std::move(groups[0]));
+  for (size_t i = 1; i < groups.size(); ++i) {
+    bool prevHasMatmul = groupHasMatmul(merged.back());
+    bool curHasMatmul = groupHasMatmul(groups[i]);
+    if (prevHasMatmul == curHasMatmul) {
+      auto &prev = merged.back();
+      prev.ops.append(groups[i].ops.begin(), groups[i].ops.end());
+      prev.nestedIfs.append(groups[i].nestedIfs.begin(),
+                            groups[i].nestedIfs.end());
+    } else {
+      merged.push_back(std::move(groups[i]));
+    }
+  }
+  return merged;
+}
+
 static inline void dumpCandidate(CandidateIf &candidate) {
   LDBG("Processing: " << candidate.ifOp);
   LDBG("  selfBlockId=" << candidate.selfBlockId);
@@ -296,15 +379,21 @@ static CandidateIf getCandidate(scf::IfOp ifOp) {
   auto selfBlockId = CVPipeline::getOpBlockId(ifOp);
   cand.selfBlockId = selfBlockId.value_or(-1);
 
-  // Group then region
-  cand.thenGroups = groupOpsInBlock(*ifOp.thenBlock());
+  // Group then region, merge consecutive same-core-type groups
+  cand.thenGroups =
+      mergeConsecutiveSameCoreType(groupOpsInBlock(*ifOp.thenBlock()));
 
-  // Group else region
+  // Group else region, merge consecutive same-core-type groups
   Block *elseBlk = ifOp.elseBlock();
   if (elseBlk) {
-    cand.elseGroups = groupOpsInBlock(*elseBlk);
+    cand.elseGroups = mergeConsecutiveSameCoreType(groupOpsInBlock(*elseBlk));
   }
 
+  // Pre-compute split decisions: only split when groups contain both
+  // matmul (CUBE) and non-matmul (VECTOR) computation.
+  cand.shouldSplitThen = hasMixedCoreTypes(cand.thenGroups);
+  cand.shouldSplitElse = hasMixedCoreTypes(cand.elseGroups);
+
   return cand;
 }
 
@@ -370,7 +459,7 @@ static void preprocessScalarDependencies(CandidateIf &cand) {
 // Part4 consumes the resulting YieldAugmentation plan.
 static void dumpYieldAugmentation(const CandidateIf &c) {
   auto &ya = c.yieldAug;
-  bool splitThen = c.thenGroups.size() >= 2;
+  bool splitThen = c.shouldSplitThen;
   auto &groups = splitThen ? c.thenGroups : c.elseGroups;
   const char *region = splitThen ? "then" : "else";
 
@@ -719,11 +808,11 @@ static void analyzeDependencies(CandidateIf &candidate) {
   // Step 2.3: Plan yield augmentation for the active region
   // Groups are in natural discovery order (block_ids appear in
   // dependency order within a sequential basic block).
-  bool splitThen = candidate.thenGroups.size() >= 2;
+  bool splitThen = candidate.shouldSplitThen;
   if (splitThen) {
     planYield(candidate, /*splitThen=*/true, ArrayRef(candidate.thenGroups),
               opToThenGroup, thenValueMap);
-  } else if (candidate.elseGroups.size() >= 2) {
+  } else if (candidate.shouldSplitElse) {
     planYield(candidate, /*splitThen=*/false, ArrayRef(candidate.elseGroups),
               opToElseGroup, elseValueMap);
   }
@@ -1409,7 +1498,7 @@ materializeCandidate(CandidateIf &c, CVPipeline::ComputeBlockIdManager &bm) {
 
   LDBG("[Part3] enter materializeCandidate hasYield=" << c.hasYield);
 
-  bool splitThen = c.thenGroups.size() >= 2;
+  bool splitThen = c.shouldSplitThen;
   auto &groups = splitThen ? c.thenGroups : c.elseGroups;
   unsigned nGroups = groups.size();
   if (nGroups < 2) {
@@ -1677,7 +1766,6 @@ void SplitIfByBlockIdPass::runOnOperation() {
   auto mainRes = walkMainLoop(module, [&](Operation *op) {
     auto funcOp = op->getParentOfType<func::FuncOp>();
     if (funcOp && (funcOp.getSymName() == kSkippedDeltaformerKernel ||
-                   funcOp.getSymName() == kSkippedChunkwiseKernel ||
                    funcOp.getSymName() == kSkippedParallelNsaFwdKernel)) {
       LDBG("Skip kernel: " << funcOp.getSymName());
       return llvm::success();

@@ -0,0 +1,72 @@
+// RUN: triton-opt --split-input-file --split-if-by-block-id %s | FileCheck %s
+
+// VCV: then side has VECTOR -> CUBE(matmul) -> VECTOR, should split.
+// CHECK-LABEL: func.func @split_then_vcv
+// CHECK: scf.if
+// CHECK: arith.addf {{.*}}ssbuffer.block_id = 94
+// CHECK: scf.if
+// CHECK: linalg.matmul {{.*}}ssbuffer.block_id = 95
+// CHECK: scf.if
+// CHECK: arith.mulf {{.*}}ssbuffer.block_id = 96
+func.func @split_then_vcv(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>, %c: f32, %d: f32, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  scf.for %iv = %c0 to %c1 step %c1 {
+    %cube = arith.addf %c, %d {ssbuffer.block_id = 0 : i32, ssbuffer.core_type = "CUBE"} : f32
+    scf.if %cond {
+      %v1 = arith.addf %c, %d {ssbuffer.block_id = 94 : i32, ssbuffer.core_type = "VECTOR"} : f32
+      %m = linalg.matmul ins(%a, %b : tensor<2x2xf32>, tensor<2x2xf32>) outs(%a : tensor<2x2xf32>) {ssbuffer.block_id = 95 : i32, ssbuffer.core_type = "CUBE"} -> tensor<2x2xf32>
+      %v2 = arith.mulf %c, %d {ssbuffer.block_id = 96 : i32, ssbuffer.core_type = "VECTOR"} : f32
+    }
+  }
+  return
+}
+
+// -----
+
+// CVC: else side has CUBE(matmul) -> VECTOR -> CUBE(matmul), should split.
+// Split ifs are re-ordered by memory dependencies (matmuls alias on %a),
+// final order is 97 -> 96 -> 95.
+// CHECK-LABEL: func.func @split_else_cvc
+// CHECK: scf.if
+// CHECK: linalg.matmul {{.*}}ssbuffer.block_id = 97
+// CHECK: scf.if
+// CHECK: arith.addf {{.*}}ssbuffer.block_id = 96
+// CHECK: scf.if
+// CHECK: linalg.matmul {{.*}}ssbuffer.block_id = 95
+func.func @split_else_cvc(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>, %c: f32, %d: f32, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  scf.for %iv = %c0 to %c1 step %c1 {
+    %cube = arith.addf %c, %d {ssbuffer.block_id = 0 : i32, ssbuffer.core_type = "CUBE"} : f32
+    scf.if %cond {
+      %x = arith.addf %c, %d {ssbuffer.block_id = 94 : i32, ssbuffer.core_type = "VECTOR"} : f32
+    } else {
+      %m1 = linalg.matmul ins(%a, %b : tensor<2x2xf32>, tensor<2x2xf32>) outs(%a : tensor<2x2xf32>) {ssbuffer.block_id = 95 : i32, ssbuffer.core_type = "CUBE"} -> tensor<2x2xf32>
+      %v1 = arith.addf %c, %d {ssbuffer.block_id = 96 : i32, ssbuffer.core_type = "VECTOR"} : f32
+      %m2 = linalg.matmul ins(%a, %b : tensor<2x2xf32>, tensor<2x2xf32>) outs(%a : tensor<2x2xf32>) {ssbuffer.block_id = 97 : i32, ssbuffer.core_type = "CUBE"} -> tensor<2x2xf32>
+    }
+  }
+  return
+}
+
+// -----
+
+// Non-CVC/VCV: else side has only 2 groups (CUBE+VECTOR), not alternating 3 groups, should NOT split.
+// CHECK-LABEL: func.func @skip_two_groups
+// CHECK-COUNT-1: scf.if
+// CHECK: arith.addf {{.*}}ssbuffer.block_id = 94
+// CHECK: linalg.matmul {{.*}}ssbuffer.block_id = 95
+func.func @skip_two_groups(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>, %c: f32, %d: f32, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  scf.for %iv = %c0 to %c1 step %c1 {
+    %cube = arith.addf %c, %d {ssbuffer.block_id = 0 : i32, ssbuffer.core_type = "CUBE"} : f32
+    scf.if %cond {
+      %x = arith.addf %c, %d {ssbuffer.block_id = 94 : i32, ssbuffer.core_type = "VECTOR"} : f32
+    } else {
+      %y = linalg.matmul ins(%a, %b : tensor<2x2xf32>, tensor<2x2xf32>) outs(%a : tensor<2x2xf32>) {ssbuffer.block_id = 95 : i32, ssbuffer.core_type = "CUBE"} -> tensor<2x2xf32>
+    }
+  }
+  return
+}