From 2e4ab46241675037d418f233f5e544bf9c0caf98 Mon Sep 17 00:00:00 2001 From: open_zk <2086579770@qq.com> Date: Tue, 1 Sep 2026 16:50:23 +0800 Subject: [PATCH] [SplitIfByBlockId] split scf.if into 3 ifs for CVC/VCV mixed core-type groups Add hasMixedCoreTypes filter so only scf.if with mixed matmul/VECTOR core-type groups (CVC or VCV) are split, merge consecutive same core-type groups to always produce a 3-way split, and remove chunkwise_fwd_kernel from the skip list. --- .../SplitIfByBlockId/SplitIfByBlockId.cpp | 118 +++++++++++++++--- .../selection_and_fallback.mlir | 72 +++++++++++ 2 files changed, 175 insertions(+), 15 deletions(-) create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/selection_and_fallback.mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/SplitIfByBlockId.cpp b/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/SplitIfByBlockId.cpp index 844210b856..115b8489c7 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/SplitIfByBlockId.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/SplitIfByBlockId.cpp @@ -59,7 +59,7 @@ #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" @@ -67,8 +67,6 @@ 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(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 &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 +mergeConsecutiveSameCoreType(SmallVector groups) { + if (groups.size() < 2) { + return groups; + } + SmallVector 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(); if (funcOp && (funcOp.getSymName() == kSkippedDeltaformerKernel || - funcOp.getSymName() == kSkippedChunkwiseKernel || funcOp.getSymName() == kSkippedParallelNsaFwdKernel)) { LDBG("Skip kernel: " << funcOp.getSymName()); return llvm::success(); diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/selection_and_fallback.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/selection_and_fallback.mlir new file mode 100644 index 0000000000..a55602820c --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/ComputeBlockOpt/SplitIfByBlockId/selection_and_fallback.mlir @@ -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 +}