diff --git a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp index 8281c0d9b6..c7d04b73ae 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp @@ -1849,7 +1849,9 @@ static void buildBeforeRegion(scf::WhileOp oldWhile, OpBuilder &bb, Location bl, bb.create(bl, condValue, carriedValues); } -// Build the after-region of the new whileOp +// Build the after-region of the new whileOp. The counter add-1 is not +// created here; it is inserted later by insertWhileCounterOps at the right +// position. counterIterArgOut is yielded as a placeholder and gets replaced. static void buildAfterRegion(scf::WhileOp oldWhile, OpBuilder &ab, Location al, ValueRange iterArgs, Value &counterIterArgOut) { Block *oldAfter = oldWhile.getAfterBody(); @@ -1871,35 +1873,12 @@ static void buildAfterRegion(scf::WhileOp oldWhile, OpBuilder &ab, Location al, if (!oldYield) return; - std::optional counterBlockId; - if (Block *doBlock = ab.getInsertionBlock()) { - for (Operation &op : llvm::reverse(*doBlock)) { - if (auto id = getOpBlockId(&op); id.has_value()) { - counterBlockId = id; - break; - } - } - } - if (!counterBlockId) - counterBlockId = getOpBlockId(oldWhile); - - Value one = ab.create(al, 1, 32); - Value nextCounter = ab.create(al, counterIterArgOut, one); - nextCounter.getDefiningOp()->setAttr(kIterCounter, ab.getUnitAttr()); - - if (counterBlockId) { - one.getDefiningOp()->setAttr(kBlockId, - ab.getI32IntegerAttr(*counterBlockId)); - nextCounter.getDefiningOp()->setAttr(kBlockId, - ab.getI32IntegerAttr(*counterBlockId)); - } - SmallVector newYieldOps; for (Value operand : oldYield->getOperands()) { Value mapped = mapper.lookupOrNull(operand); newYieldOps.push_back(mapped ? mapped : operand); } - newYieldOps.push_back(nextCounter); + newYieldOps.push_back(counterIterArgOut); ab.create(al, newYieldOps); } @@ -1952,6 +1931,77 @@ setupWhileIterArgCounter(const MainLoop &loop, OpBuilder &builder) { return {counterIterArg, newWhile}; } +// Insert the iterCounter add-1 (and its constant 1) right after the last op +// sharing the block_id of the first op that consumes mainLoop.iterCounter. +// Must run after processTensorDependencies (dispatch ops land there later). +// Fallback when no counter user exists: last op's block_id, mirroring the +// legacy buildAfterRegion placement. +static void insertWhileCounterOps(const MainLoop &mainLoop) { + if (!mainLoop.isWhile() || !mainLoop.iterCounter) + return; + + auto whileOp = cast(mainLoop.getOperation()); + Block *doBlock = whileOp.getAfterBody(); + if (!doBlock) + return; + + std::optional counterBlockId; + for (Operation &op : *doBlock) { + if (llvm::is_contained(op.getOperands(), mainLoop.iterCounter)) { + counterBlockId = getOpBlockId(&op); + break; + } + } + if (!counterBlockId) { + for (Operation &op : llvm::reverse(*doBlock)) { + if (auto id = getOpBlockId(&op); id.has_value()) { + counterBlockId = id; + break; + } + } + if (!counterBlockId) + return; + } + + Operation *lastWithBlockId = nullptr; + for (Operation &op : *doBlock) { + auto id = getOpBlockId(&op); + if (id.has_value() && *id == *counterBlockId) + lastWithBlockId = &op; + } + + Location loc = lastWithBlockId->getLoc(); + OpBuilder builder(mainLoop.getContext()); + builder.setInsertionPointAfter(lastWithBlockId); + IntegerAttr blockIdAttr = builder.getI32IntegerAttr(*counterBlockId); + + Value one = builder.create( + loc, 1, mainLoop.iterCounter.getType().getIntOrFloatBitWidth()); + one.getDefiningOp()->setAttr(kBlockId, blockIdAttr); + + Value nextCounter = + builder.create(loc, mainLoop.iterCounter, one); + Operation *iterAddOp = nextCounter.getDefiningOp(); + iterAddOp->setAttr(kBlockId, blockIdAttr); + iterAddOp->setAttr(kIterCounter, builder.getUnitAttr()); + + // Replace the placeholder counterIterArgOut operand in the yield with + // nextCounter so the iter_arg increments each iteration. + Operation *yieldOp = doBlock->getTerminator(); + SmallVector newOperands(yieldOp->getOperands().begin(), + yieldOp->getOperands().end()); + bool replaced = false; + for (Value &operand : newOperands) { + if (operand == mainLoop.iterCounter) { + operand = nextCounter; + replaced = true; + break; + } + } + if (replaced) + yieldOp->setOperands(newOperands); +} + static int addInnerMultiBuffer(MainLoop mainLoop, OpBuilder &builder, scope::ScopeOp vectorScope, int &groupId, bool &i1Found) { @@ -2087,6 +2137,10 @@ static int addInnerMultiBuffer(MainLoop mainLoop, OpBuilder &builder, return -1; } + // WhileOp only: insert the counter add-1 now that the dispatch ops that + // use mainLoop.iterCounter have been emitted. + insertWhileCounterOps(mainLoop); + LLVM_DEBUG(llvm::dbgs() << "[addInnerMultiBuffer] DONE\n"); return 0; } diff --git a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp b/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp index 956c8722a4..cc1e8009ac 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp @@ -42,8 +42,7 @@ namespace { static constexpr llvm::StringLiteral interceptrFunc[]{ "chunk_abc_bwd_kernel_dh", "flash_varlen_fwd_kernel", - "chunk_gsa_bwd_k_kernel_dqkvg", "_jagged_flash_attention_bwd_basic_kernel", - "_sparse_decode_kernel", "chunk_gsa_fwd_k_kernel_intra", + "chunk_gsa_bwd_k_kernel_dqkvg", "_sparse_decode_kernel", "_sparse_decode_model1_kernel", "sparse_flash_attention_grad_kernel"}; static LogicalResult verifyFuncNames(ModuleOp module) { diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir index bbd43eb386..eda2f3cf85 100644 --- a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir @@ -35,6 +35,9 @@ // (3 block-args instead of 2; the new arg is i32). // - arith.addi counter, 1 is present in the do-region for the yield. // - Multi-buffer producer/consumer with scf.if dispatch works as forOp. +// - The counter add-1 op (and its constant-1) is relocated to the +// block_id of the first op that consumes the counter arg (here: the +// producer scf.if at block_id = 7), not its old fallback position. // CHECK-LABEL: func.func @test_while_mainloop_bufnum_two // Two UB allocs (ping/pong). @@ -42,14 +45,14 @@ // CHECK-DAG: memref.alloc() : memref<128xf32, #hivm.address_space> // WhileOp's do-region bb0 has 3 block-args now (the new i32 counter is the last one). // CHECK: ^bb0(%{{.*}}: tensor<128xf32>, %{{.*}}: i32, %{{.*}}: i32): -// Producer scf.if dispatch (the original counter increment lives inside this region). +// Producer scf.if dispatch. // CHECK: scf.if // CHECK: hivm.hir.copy +// arith.addi increment for the multi-buffer counter (block_id=7, the first user's block). +// CHECK: %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 7 : i32, ssbuffer.iterCounter} : i32 // Consumer scf.if dispatch returning tensor. // CHECK: scf.if {{.*}} -> (tensor<128xf32>) // CHECK: bufferization.to_tensor -// arith.addi increment for the multi-buffer counter is present (block_id=10), tagged iterCounter. -// CHECK: %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 10 : i32, ssbuffer.iterCounter} : i32 // Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop. // CHECK: } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64 @@ -63,11 +66,11 @@ // Producer alloc + copy before the whileOp. // CHECK-DAG: memref.alloc() : memref<64xf16, #hivm.address_space> // CHECK-DAG: memref.alloc() : memref<64xf16, #hivm.address_space> +// arith.addi increment for the multi-buffer counter (block_id=7, the first user's block). +// CHECK: %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 7 : i32, ssbuffer.iterCounter} : i32 // Consumer scf.if + to_tensor inside the do-region. // CHECK: scf.if {{.*}} -> (tensor<64xf16>) // CHECK: bufferization.to_tensor -// arith.addi increment for the multi-buffer counter is present (block_id=12), tagged iterCounter. -// CHECK: %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 12 : i32, ssbuffer.iterCounter} : i32 // Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop. // CHECK: } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64 @@ -90,8 +93,8 @@ // Producer-side dispatch. // CHECK: scf.if // CHECK: hivm.hir.copy -// arith.addi increment for the multi-buffer counter is present (block_id=12), tagged iterCounter. -// CHECK: %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 12 : i32, ssbuffer.iterCounter} : i32 +// arith.addi increment for the multi-buffer counter (block_id=8, the first user's block). +// CHECK: %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 8 : i32, ssbuffer.iterCounter} : i32 // Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop. // CHECK: } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64 @@ -121,6 +124,34 @@ // Not implemented in this initial drop — to be added when the negative case // is needed for a regression guard. +// T-while-G: Regression guard for relocateWhileCounterOps. +// The counter add-1 (and its constant 1) used to be appended at the end of +// the do-region, carrying the whileOp's own block_id (= 99 here, distinct +// from both producer and consumer). The pass now relocates them to the +// block_id of the first op that consumes the counter iter_arg — the +// producer scf.if at block_id = 5. +// We assert: +// - The counter add-1 carries block_id = 5 (NOT 99, NOT 50). +// - The constant 1 used by the addi also carries block_id = 5. +// - The scf.yield that consumes the add-1's SSA value still references +// // the same %result (SSA is preserved across the move). + +// CHECK-LABEL: func.func @test_while_counter_relocation +// Multi-buffer producer (two allocs at block 5) emits the dispatch. +// CHECK-DAG: memref.alloc() : memref<16xf32, #hivm.address_space> +// CHECK-DAG: memref.alloc() : memref<16xf32, #hivm.address_space> +// CHECK: scf.if +// CHECK: hivm.hir.copy +// Constant 1 carrying block_id = 5 (relocated with the addi). +// CHECK: %{{.+}} = arith.constant {ssbuffer.block_id = 5 : i32} 1 : i32 +// Counter add-1 carries block_id = 5 and is tagged iterCounter. +// CHECK: %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 5 : i32, ssbuffer.iterCounter} : i32 +// scf.yield forwards the counter (referenced by the addi's %addi-res), proving +// SSA is preserved. +// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.+}} +// Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop. +// CHECK: } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64 + // ---- Inputs ---- module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">, @@ -277,4 +308,37 @@ module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">, } {hivm.tcore_type = #hivm.tcore_type} return } + + // T-while-G: dedicated regression test for relocateWhileCounterOps. + // Producer at block_id = 5, consumer at block_id = 50, whileOp at block_id = 99. + // Before the fix: counter add-1 would land at block_id = 99 (whileOp's own id). + // After the fix: counter add-1 (and its constant 1) land at block_id = 5 + // (the block_id of the first op that consumes the counter iter_arg). + func.func @test_while_counter_relocation() { + %c0_i32 = arith.constant 0 : i32 + %c10_i32 = arith.constant 10 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst_zero = arith.constant 0.0 : f32 + %init = tensor.empty() : tensor<16xf32> + %carry = linalg.fill ins(%cst_zero : f32) outs(%init : tensor<16xf32>) -> tensor<16xf32> + scope.scope : () -> () { + %result:2 = scf.while (%arg0 = %carry, %arg1 = %c0_i32) + : (tensor<16xf32>, i32) -> (tensor<16xf32>, i32) { + %cmp = arith.cmpi slt, %arg1, %c10_i32 {ssbuffer.block_id = 16 : i32} : i32 + scf.condition(%cmp) %arg0, %arg1 : tensor<16xf32>, i32 + } do { + ^bb0(%arg0: tensor<16xf32>, %arg1: i32): + // Producer block_id = 5. + %alloc = memref.alloc() {ssbuffer.block_id = 5 : i32} : memref<16xf32> + %prod = bufferization.to_tensor %alloc {ssbuffer.block_id = 5 : i32} : memref<16xf32> to tensor<16xf32> + // Consumer block_id = 50 (very different from 5, to make the relocation + // verdict unambiguous). + %consumed = arith.addf %prod, %prod {ssbuffer.block_id = 50 : i32} : tensor<16xf32> + %next = arith.addi %arg1, %c1_i32 : i32 + scf.yield %consumed, %next : tensor<16xf32>, i32 + } attributes {ssbuffer.main_loop = 1 : i64, ssbuffer.block_id = 99 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } }