diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index b115f4a901..8a5eec116f 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -103,6 +103,7 @@ std::unique_ptr createPTOOutlineSIMTSectionsPass(); std::unique_ptr createPTOInferVPTOVecScopePass(); std::unique_ptr createVPTOExpandWrapperOpsPass(); std::unique_ptr createVPTOSoftPostUpdatePass(); +std::unique_ptr createVPTOGuardedLICMPass(); std::unique_ptr createPTOPrintAddressAnalysisPass(); std::unique_ptr createPTOVPTOPtrBoundaryPass(); std::unique_ptr diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index ca5e10b4f0..528b17d4e8 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -1340,6 +1340,34 @@ def VPTOSoftPostUpdate "mlir::scf::SCFDialect"]; } +def VPTOGuardedLICM : Pass<"vpto-guarded-licm", "func::FuncOp"> { + let summary = "Hoist loop-invariant guarded expressions out of scf.if regions"; + let description = [{ + The generic loop-invariant code motion pass only inspects the top-level + operations of the scf.for body. A guarded access pattern that computes a + loop-invariant base address chain inside an scf.if whose condition depends + on the induction variable is therefore never exposed to LICM, and the + surrounding scf.if cannot itself be hoisted because its region contains + side-effecting memory operations. + + This pass extracts the pure, speculatable scalar and address + subexpressions (integer, index, and pto.ptr results) of such guarded + regions and moves them in dependency (topological) order to just before + the enclosing scf.for. IV-dependent arithmetic, side-effecting + operations, div/rem with unknown divisors, vector/mask/tile computations, + and operations inside nested scf.for regions are left in place. Loops + are processed innermost-first so invariant chains can climb out through + arbitrarily deep guards and nested loops. When a property cannot be + proven the pass conservatively keeps the expression in place. The + rewrite is enabled only for the A5 VPTO backend. + }]; + let constructor = "mlir::pto::createVPTOGuardedLICMPass()"; + let dependentDialects = ["mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::arith::ArithDialect", + "mlir::scf::SCFDialect"]; +} + def PTOPrintAddressAnalysis : Pass<"pto-print-address-analysis", "func::FuncOp"> { let summary = "Print cached PTO value-evolution and VPTO address facts"; diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 05067d9e7c..04acf38839 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -49,6 +49,7 @@ add_mlir_dialect_library(PTOTransforms VPTOMaskSimplify.cpp VPTOExpandWrapperOps.cpp VPTOSoftPostUpdate.cpp + VPTOGuardedLICM.cpp PTOPrintAddressAnalysis.cpp VPTOScheduler/VPTORegPressureTracker.cpp VPTOScheduler/VPTOSchedBoundary.cpp diff --git a/lib/PTO/Transforms/VPTOGuardedLICM.cpp b/lib/PTO/Transforms/VPTOGuardedLICM.cpp new file mode 100644 index 0000000000..ca6f1b24b6 --- /dev/null +++ b/lib/PTO/Transforms/VPTOGuardedLICM.cpp @@ -0,0 +1,283 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VPTOGuardedLICM.cpp -----------------------------------------------===// +// +// Hoist loop-invariant scalar and address expressions out of guarded scf.if +// regions nested inside scf.for loops. +// +// The generic MLIR loop-invariant code motion pass only enqueues the +// top-level operations of the scf.for body. For a guarded access pattern +// +// scf.for %iv = ... +// %in_bounds = ... // depends on the IV +// scf.if %in_bounds { +// %base = ... // loop-invariant base address chain +// %dynamic = ... %iv ... // IV-dependent offset +// pto.store ... +// } +// +// the invariant chain inside the guard is therefore never seen by LICM, and +// the whole scf.if cannot be hoisted either: its condition depends on the IV +// and its region contains side-effecting memory operations. This pass +// extracts only the safe, speculatable scalar/address subexpressions of the +// guard and moves them in dependency (topological) order to just before the +// loop, leaving IV-dependent arithmetic, side-effecting operations and +// vector/container computations in place. +// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/STLExtras.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VPTOGUARDEDLICM +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; + +namespace { + +// Address-like scalar result types we are willing to hoist: plain integers, +// index values and pto.ptr addresses. Vector registers, masks, tile +// objects and other container types are deliberately excluded so that the +// hoisted code cannot grow vector register pressure outside the guarded +// region. +static bool isHoistableScalarType(Type type) { + return isa(type); +} + +// The whole integer division/remainder family is never hoisted, including the +// "total" variants (arith.floordivsi/ceildivsi/ceildivui) that do not even +// carry ConditionallySpeculatable: dividing by zero or INT_MIN / -1 is +// undefined behavior, so evaluating them before the guard could introduce UB +// that the original program never executed. divf/remf are not hoisted +// either: a float div/rem can raise IEEE exceptions or trap, and the generic +// pure check does not protect the short-circuiting case. +static bool isDivOrRem(Operation *op) { + return isa(op); +} + +// Explicit allow-list of deterministic, referentially transparent scalar +// operations. The generic isPure() is intentionally not the gate here: +// several PTO operations are marked pure but are not referentially +// transparent, e.g. pto.get_clock32/64 (time sampling), pto.get_vms4_sr +// (VMS4 status register), and pto.vote_*/pto.shuffle_*/pto.redux_* +// (active-lane collectives). Hoisting any of them changes the value the +// guard observes, so only the operations listed below may move out of a +// guard; everything else stays in place. Both the signed and the unsigned +// index<->integer casts (arith.index_cast / arith.index_castui) are allowed: +// VPTO soft post-update emits the unsigned form for offset/block computations, +// and excluding it would leave the whole guarded base chain behind. +static bool isHoistableScalarOp(Operation *op) { + return isa(op); +} + +// An operation is a hoist candidate only when it is on the deterministic +// allow-list, still side-effect free (defensive double check), region-free, +// produces only address-like scalars, and is not a terminator. +static bool isHoistCandidate(Operation *op) { + if (op->hasTrait()) { + return false; + } + if (!op->getRegions().empty()) { + return false; + } + if (!isHoistableScalarOp(op)) { + return false; + } + if (!isPure(op)) { + return false; + } + if (isDivOrRem(op)) { + return false; + } + for (Type resultType : op->getResultTypes()) { + if (!isHoistableScalarType(resultType)) { + return false; + } + } + return true; +} + +// Whether a value is defined outside the loop, i.e. available before the loop +// starts. Loop induction variables, iter_args of the loop body and every +// value produced inside the loop fail this check. +static bool isDefinedOutsideLoop(Value value, Operation *loopOp) { + Operation *holder = nullptr; + if (auto blockArg = dyn_cast(value)) { + holder = blockArg.getOwner()->getParentOp(); + } else { + holder = value.getDefiningOp(); + } + if (!holder) { + return true; + } + Operation *ancestor = holder; + while (ancestor) { + if (ancestor == loopOp) { + return false; + } + ancestor = ancestor->getParentOp(); + } + return true; +} + +// Number of enclosing scf.for loops; used to sort loops innermost-first. +static int loopNestingDepth(scf::ForOp forOp) { + int depth = 0; + Operation *op = forOp.getOperation(); + while (op != nullptr) { + if (isa(op)) { + ++depth; + } + op = op->getParentOp(); + } + return depth; +} + +// Extract every hoistable invariant subexpression that lives inside an scf.if +// region nested in the loop (and not inside another scf.for), moving the +// expressions in dependency order to just before the loop. +static void hoistInvariantsFromGuards(scf::ForOp forOp) { + // Collect candidates. walk() covers every nested scf.if region; the + // innermost scf.for ancestor test keeps nested loops in charge of their own + // decisions (they are processed before this loop because loops are visited + // innermost-first). + SmallVector candidates; + forOp->walk([&](Operation *op) { + if (op == forOp.getOperation()) { + return; + } + scf::IfOp enclosingIf = op->getParentOfType(); + if (enclosingIf == nullptr) { + return; // not inside a guarded region + } + scf::ForOp enclosingFor = op->getParentOfType(); + if (enclosingFor != forOp) { + return; // inside a nested loop; that loop owns the decision + } + if (isHoistCandidate(op)) { + candidates.push_back(op); + } + }); + + if (candidates.empty()) { + return; + } + + // Iterate to a fixed point: an op joins the hoist set only when all of its + // operands are either defined outside the loop or produced by already-hoisted + // ops. The join order is therefore a valid topological order. + DenseSet available; + DenseSet hoistedSet; + SmallVector hoisted; + bool changed = true; + while (changed) { + changed = false; + for (Operation *op : candidates) { + if (hoistedSet.count(op)) { + continue; + } + bool allOperandsAvailable = true; + for (Value operand : op->getOperands()) { + bool operandAvailable = available.count(operand) != 0; + bool operandFromOutside = isDefinedOutsideLoop(operand, forOp); + if (!operandAvailable && !operandFromOutside) { + allOperandsAvailable = false; + break; + } + } + if (!allOperandsAvailable) { + continue; + } + hoisted.push_back(op); + hoistedSet.insert(op); + for (Value result : op->getResults()) { + available.insert(result); + } + changed = true; + } + } + + for (Operation *op : hoisted) { + op->moveBefore(forOp); + } +} + +// The rewrite is restricted to the A5 VPTO module; keep other targets (A3, +// EmitC) bit-for-bit identical until they have their own performance +// validation. +static bool isA5VPTOModule(func::FuncOp func) { + ModuleOp module = func->getParentOfType(); + while (module != nullptr) { + auto arch = module->getAttrOfType("pto.target_arch"); + if (arch) { + return arch.getValue() == "a5"; + } + module = module->getParentOfType(); + } + return false; +} + +struct VPTOGuardedLICM + : public pto::impl::VPTOGuardedLICMBase { + using pto::impl::VPTOGuardedLICMBase::VPTOGuardedLICMBase; + + void runOnOperation() override { + func::FuncOp func = getOperation(); + + if (!isA5VPTOModule(func)) { + return; + } + + // Process loops innermost-first so invariants can climb out through nested + // guards and nested loops: after an inner loop hoists a chain to its own + // preheader (still inside the guard of an outer loop), the outer loop pass + // re-collects the same chain and hoists it once more. + SmallVector loops; + func.walk([&](scf::ForOp forOp) { loops.push_back(forOp); }); + llvm::sort(loops, [](scf::ForOp lhs, scf::ForOp rhs) { + return loopNestingDepth(lhs) > loopNestingDepth(rhs); + }); + + for (scf::ForOp forOp : loops) { + hoistInvariantsFromGuards(forOp); + } + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVPTOGuardedLICMPass() { + return std::make_unique(); +} diff --git a/test/lit/vpto/vpto_guarded_licm.pto b/test/lit/vpto/vpto_guarded_licm.pto new file mode 100644 index 0000000000..e67a79f8fc --- /dev/null +++ b/test/lit/vpto/vpto_guarded_licm.pto @@ -0,0 +1,346 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// The generic MLIR LICM only inspects top-level scf.for body operations, so +// a loop-invariant base address chain computed inside an scf.if guard is +// never hoisted. VPTOGuardedLICM extracts pure, speculatable scalar/address +// subexpressions of the guard and moves them in dependency order in front of +// the loop, while IV-dependent arithmetic, side-effecting operations and +// non-scalar (vector/float) computations stay inside the guard. +// +// RUN: pto-test-opt --vpto-guarded-licm %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + + // The base chain (extsi %lane -> muli -> addi) is loop-invariant and must + // be hoisted in front of scf.for; the IV-dependent offset chain and the + // guarded load/store must remain inside the scf.if. The f32 addf result + // is not an address-like scalar and must stay inside the guard as well. + func.func @guarded_licm_basic(%buf: !pto.ptr, %lane: i32, + %count: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1000 = arith.constant 1000 : index + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + %c65536_i64 = arith.constant 65536 : i64 + scf.for %iv = %c0 to %count step %c1 { + %in_bounds = arith.cmpi slt, %iv, %c1000 : index + scf.if %in_bounds { + %w64 = arith.extsi %lane : i32 to i64 + %slot = arith.muli %w64, %c32_i64 : i64 + %range_base = arith.addi %slot, %c65536_i64 : i64 + %iv64 = arith.index_cast %iv : index to i64 + %dynamic = arith.muli %iv64, %c128_i64 : i64 + %offset64 = arith.addi %range_base, %dynamic : i64 + %addr = arith.index_cast %offset64 : i64 to index + %v = pto.load_scalar %buf[%addr] : !pto.ptr -> f32 + %v2 = arith.addf %v, %v : f32 + pto.store_scalar %v2, %buf[%addr] : !pto.ptr, f32 + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_basic + // CHECK: %[[W64:.*]] = arith.extsi %[[LANE:.*]] : i32 to i64 + // CHECK: %[[SLOT:.*]] = arith.muli %[[W64]], %{{.*}} : i64 + // CHECK: %[[RANGEBASE:.*]] = arith.addi %[[SLOT]], %{{.*}} : i64 + // CHECK-NOT: arith.index_cast + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: %[[IV64:.*]] = arith.index_cast %[[IV:.*]] : index to i64 + // CHECK: %[[DYN:.*]] = arith.muli %[[IV64]], %{{.*}} : i64 + // CHECK: %[[OFF:.*]] = arith.addi %[[RANGEBASE]], %[[DYN]] : i64 + // CHECK: %[[ADDR:.*]] = arith.index_cast %[[OFF]] : i64 to index + // CHECK: pto.load_scalar + // CHECK: arith.addf + // CHECK: pto.store_scalar + + // An unsigned index<->integer cast (arith.index_castui) must also be + // hoistable: VPTO soft post-update emits this form for unsigned offset and + // block computations, so it can appear in the guarded invariant base chain. + // The invariant chain (including the index_castui) moves in front of the + // loop; the IV-dependent unsigned offset chain stays inside the guard. + func.func @guarded_licm_index_castui(%buf: !pto.ptr, %lane: i32, + %count: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1000 = arith.constant 1000 : index + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + %c65536_i64 = arith.constant 65536 : i64 + %c0_i32 = arith.constant 0 : i32 + scf.for %iv = %c0 to %count step %c1 { + %in_bounds = arith.cmpi slt, %iv, %c1000 : index + scf.if %in_bounds { + %w64 = arith.extsi %lane : i32 to i64 + %slot = arith.muli %w64, %c32_i64 : i64 + %base64 = arith.addi %slot, %c65536_i64 : i64 + %base = arith.index_castui %base64 : i64 to index + %iv64 = arith.index_castui %iv : index to i64 + %dynamic = arith.muli %iv64, %c128_i64 : i64 + %dyn_index = arith.index_cast %dynamic : i64 to index + %addr = arith.addi %base, %dyn_index : index + pto.store_scalar %c0_i32, %buf[%addr] : !pto.ptr, i32 + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_index_castui + // CHECK: %[[W64:.*]] = arith.extsi %{{.*}} : i32 to i64 + // CHECK: %[[SLOT:.*]] = arith.muli %[[W64]], %{{.*}} : i64 + // CHECK: %[[BASE64:.*]] = arith.addi %[[SLOT]], %{{.*}} : i64 + // CHECK: %[[BASE:.*]] = arith.index_castui %[[BASE64]] : i64 to index + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: %[[IV64:.*]] = arith.index_castui %[[IV:.*]] : index to i64 + // CHECK: %[[DYN:.*]] = arith.muli %[[IV64]], %{{.*}} : i64 + // CHECK: %[[DYNIDX:.*]] = arith.index_cast %[[DYN]] : i64 to index + // CHECK: %[[ADDR:.*]] = arith.addi %[[BASE]], %[[DYNIDX]] : index + // CHECK: pto.store_scalar + + // Invariants must climb out of nested scf.if guards as well. + func.func @guarded_licm_nested_if(%buf: !pto.ptr, %lane: i32, + %count: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1000 = arith.constant 1000 : index + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + %c65536_i64 = arith.constant 65536 : i64 + scf.for %iv = %c0 to %count step %c1 { + %outer_cond = arith.cmpi slt, %iv, %c1000 : index + scf.if %outer_cond { + %inner_cond = arith.cmpi sge, %iv, %c0 : index + scf.if %inner_cond { + %w64 = arith.extsi %lane : i32 to i64 + %slot = arith.muli %w64, %c32_i64 : i64 + %range_base = arith.addi %slot, %c65536_i64 : i64 + %iv64 = arith.index_cast %iv : index to i64 + %dynamic = arith.muli %iv64, %c128_i64 : i64 + %offset64 = arith.addi %range_base, %dynamic : i64 + %addr = arith.index_cast %offset64 : i64 to index + %v = pto.load_scalar %buf[%addr] : !pto.ptr -> f32 + pto.store_scalar %v, %buf[%addr] : !pto.ptr, f32 + } + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_nested_if + // CHECK: %[[W64:.*]] = arith.extsi %[[LANE:.*]] : i32 to i64 + // CHECK: %[[SLOT:.*]] = arith.muli %[[W64]], %{{.*}} : i64 + // CHECK: %[[RANGEBASE:.*]] = arith.addi %[[SLOT]], %{{.*}} : i64 + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: %[[IV64:.*]] = arith.index_cast %[[IV:.*]] : index to i64 + // CHECK: %[[DYN:.*]] = arith.muli %[[IV64]], %{{.*}} : i64 + // CHECK: %[[OFF:.*]] = arith.addi %[[RANGEBASE]], %[[DYN]] : i64 + // CHECK: pto.load_scalar + // CHECK: pto.store_scalar + + // Innermost-first processing: the inner loop sits inside an outer guard. The + // invariant chain must climb from the inner loop's guard, out of the outer + // guard, all the way in front of the outer loop; the inner-IV-dependent + // offset stays inside the inner guard. + func.func @guarded_licm_nested_loops(%buf: !pto.ptr, %lane: i32, + %outer_count: index, %inner_count: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1000 = arith.constant 1000 : index + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + %c65536_i64 = arith.constant 65536 : i64 + scf.for %o = %c0 to %outer_count step %c1 { + %outer_cond = arith.cmpi slt, %o, %c1000 : index + scf.if %outer_cond { + scf.for %i = %c0 to %inner_count step %c1 { + %active = arith.cmpi slt, %i, %c1000 : index + scf.if %active { + %w64 = arith.extsi %lane : i32 to i64 + %slot = arith.muli %w64, %c32_i64 : i64 + %range_base = arith.addi %slot, %c65536_i64 : i64 + %iv64 = arith.index_cast %i : index to i64 + %dynamic = arith.muli %iv64, %c128_i64 : i64 + %offset64 = arith.addi %range_base, %dynamic : i64 + %addr = arith.index_cast %offset64 : i64 to index + %v = pto.load_scalar %buf[%addr] : !pto.ptr -> f32 + pto.store_scalar %v, %buf[%addr] : !pto.ptr, f32 + } + } + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_nested_loops + // CHECK: %[[W64:.*]] = arith.extsi %{{.*}} : i32 to i64 + // CHECK: %[[SLOT:.*]] = arith.muli %[[W64]], %{{.*}} : i64 + // CHECK: %[[RANGEBASE:.*]] = arith.addi %[[SLOT]], %{{.*}} : i64 + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: %[[IV64:.*]] = arith.index_cast %[[IV:.*]] : index to i64 + // CHECK: %[[DYN:.*]] = arith.muli %[[IV64]], %{{.*}} : i64 + // CHECK: %[[OFF:.*]] = arith.addi %[[RANGEBASE]], %[[DYN]] : i64 + // CHECK: pto.load_scalar + // CHECK: pto.store_scalar + + // Negative (IV-dependent): every address expression transitively depends + // on the induction variable, so nothing may be hoisted in front of the loop. + func.func @guarded_licm_negative_iv(%buf: !pto.ptr, + %count: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1000 = arith.constant 1000 : index + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + scf.for %iv = %c0 to %count step %c1 { + %in_bounds = arith.cmpi slt, %iv, %c1000 : index + scf.if %in_bounds { + %iv64 = arith.index_cast %iv : index to i64 + %slot = arith.muli %iv64, %c32_i64 : i64 + %dynamic = arith.muli %iv64, %c128_i64 : i64 + %combined = arith.addi %slot, %dynamic : i64 + %addr = arith.index_cast %combined : i64 to index + %v = pto.load_scalar %buf[%addr] : !pto.ptr -> f32 + pto.store_scalar %v, %buf[%addr] : !pto.ptr, f32 + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_negative_iv + // CHECK-NOT: arith.index_cast + // CHECK-NOT: arith.muli + // CHECK-NOT: arith.addi + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: arith.index_cast + // CHECK: arith.muli + // CHECK: arith.addi + // CHECK: pto.load_scalar + // CHECK: pto.store_scalar + + // Negative (non-speculatable): the divsi divisor is an unknown function + // argument, so even though its operands are loop-invariant and hoistable, + // the division itself must stay inside the guard. + func.func @guarded_licm_negative_div(%buf: !pto.ptr, %lane: i32, + %count: index, %divisor: i64) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1000 = arith.constant 1000 : index + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + scf.for %iv = %c0 to %count step %c1 { + %in_bounds = arith.cmpi slt, %iv, %c1000 : index + scf.if %in_bounds { + %w64 = arith.extsi %lane : i32 to i64 + %slot = arith.muli %w64, %c32_i64 : i64 + %quotient = arith.divsi %slot, %divisor : i64 + %addr = arith.index_cast %slot : i64 to index + %v = pto.load_scalar %buf[%addr] : !pto.ptr -> f32 + pto.store_scalar %v, %buf[%addr] : !pto.ptr, f32 + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_negative_div + // CHECK: %[[W64:.*]] = arith.extsi %[[LANE:.*]] : i32 to i64 + // CHECK: %[[SLOT:.*]] = arith.muli %[[W64]], %{{.*}} : i64 + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: %[[QUOTIENT:.*]] = arith.divsi %[[SLOT]], %{{.*}} : i64 + // CHECK: pto.load_scalar + // CHECK: pto.store_scalar + + // Negative (SIMT referential transparency): pto.get_clock32 reads the + // clock, pto.get_vms4_sr reads the VMS4 status register, and + // pto.vote_ballot / pto.shuffle_idx are active-lane collectives. Although + // these ops are marked pure, hoisting any of them out of the guard changes + // the value observed inside the guard, so none may move even though every + // operand is defined outside the loop. + func.func @guarded_licm_negative_simt(%buf: !pto.ptr, %v: i32, + %count: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1_i32 = arith.constant 1 : i32 + %c1000 = arith.constant 1000 : index + %pred = arith.cmpi eq, %v, %v : i32 + scf.for %iv = %c0 to %count step %c1 { + %in_bounds = arith.cmpi slt, %iv, %c1000 : index + scf.if %in_bounds { + %clock = pto.get_clock32 : i32 + %s0, %s1, %s2, %s3 = pto.get_vms4_sr : i16, i16, i16, i16 + %ballot = pto.vote_ballot %pred : i1 -> i32 + %shuf = pto.shuffle_idx %v, %c1_i32 : i32, i32 -> i32 + %c0_i32 = arith.constant 0 : i32 + pto.store_scalar %c0_i32, %buf[%iv] : !pto.ptr, i32 + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_negative_simt + // CHECK-NOT: pto.get_clock32 + // CHECK-NOT: pto.get_vms4_sr + // CHECK-NOT: pto.vote_ballot + // CHECK-NOT: pto.shuffle_idx + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: pto.get_clock32 + // CHECK: pto.get_vms4_sr + // CHECK: pto.vote_ballot + // CHECK: pto.shuffle_idx + // CHECK: pto.store_scalar + + // Negative (non-speculatable "total" div): arith.floordivsi is in the total + // integer-binary family without ConditionallySpeculatable, yet dividing by + // zero (or INT_MIN / -1) is still undefined behavior. Even though its + // operands are loop-invariant, the division must stay inside the guard. + func.func @guarded_licm_negative_floordiv(%buf: !pto.ptr, + %count: index, %x: i64, + %divisor: i64) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1000 = arith.constant 1000 : index + %c0_i32 = arith.constant 0 : i32 + scf.for %iv = %c0 to %count step %c1 { + %in_bounds = arith.cmpi slt, %iv, %c1000 : index + scf.if %in_bounds { + %quotient = arith.floordivsi %x, %divisor : i64 + %addr = arith.index_cast %quotient : i64 to index + pto.store_scalar %c0_i32, %buf[%addr] : !pto.ptr, i32 + } + } + return + } + + // CHECK-LABEL: func.func @guarded_licm_negative_floordiv + // CHECK-NOT: arith.floordivsi + // CHECK: scf.for + // CHECK: arith.cmpi + // CHECK: scf.if + // CHECK: %[[QUOTIENT:.*]] = arith.floordivsi %{{.*}}, %{{.*}} : i64 + // CHECK: pto.store_scalar +} diff --git a/test/lit/vpto/vpto_guarded_licm_issue1328.pto b/test/lit/vpto/vpto_guarded_licm_issue1328.pto new file mode 100644 index 0000000000..1ef10cba9d --- /dev/null +++ b/test/lit/vpto/vpto_guarded_licm_issue1328.pto @@ -0,0 +1,100 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// End-to-end regression for issue #1328: a pto.section.simt body outlines into +// a helper whose scf.if guard computes a long loop-invariant base-address +// chain. The full VPTO emission pipeline (VPTOGuardedLICM between soft +// post-update and the generic LICM) must hoist the whole invariant chain in +// front of the scf.for, leaving only the IV-dependent i*32 portion and the +// guarded store inside the guard. +// +// RUN: ptoas --pto-backend=vpto --pto-arch=a5 --emit-vpto %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @simt_guarded_licm( + %dst: !pto.ptr, + %w: i32, + %valid_end: i32, + %local_end: i32, + %total_start: i32, + %trip_count: index) attributes {pto.entry} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c31_i32 = arith.constant 31 : i32 + %c32_i32 = arith.constant 32 : i32 + %c3_i64 = arith.constant 3 : i64 + %c31_i64 = arith.constant 31 : i64 + %c32_i64 = arith.constant 32 : i64 + %c128_i64 = arith.constant 128 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %minus_one = arith.constant -1 : i32 + + pto.section.simt<<<512, 1, 1>>> { + %tid = pto.get_tid_x : i32 + %lane = arith.andi %tid, %c31_i32 : i32 + scf.for %i = %c0 to %trip_count step %c1 { + %i32 = arith.index_cast %i : index to i32 + %iter_offset = arith.muli %i32, %c32_i32 : i32 + %logical_index = arith.addi %iter_offset, %lane : i32 + %in_bounds = arith.cmpi slt, %logical_index, %local_end : i32 + scf.if %in_bounds { + %w64 = arith.extsi %w : i32 to i64 + %slot = arith.andi %w64, %c3_i64 : i64 + %slot_base = arith.muli %slot, %c1024_i64 : i64 + %valid64 = arith.extsi %valid_end : i32 to i64 + %local64 = arith.extsi %local_end : i32 to i64 + %range_base = arith.minsi %valid64, %local64 : i64 + %tid64 = arith.extsi %tid : i32 to i64 + %lane64 = arith.andi %tid64, %c31_i64 : i64 + %total_start64 = arith.extsi %total_start : i32 to i64 + %base0 = arith.addi %slot_base, %range_base : i64 + %base1 = arith.addi %base0, %lane64 : i64 + %base2 = arith.addi %base1, %c128_i64 : i64 + %base = arith.subi %base2, %total_start64 : i64 + + %i64 = arith.extsi %i32 : i32 to i64 + %dynamic_offset = arith.muli %i64, %c32_i64 : i64 + %offset = arith.addi %base, %dynamic_offset : i64 + %index = arith.index_cast %offset : i64 to index + pto.store %minus_one, %dst[%index] : !pto.ptr, i32 + } + } + } + return + } +} + +// CHECK-LABEL: func.func private @simt_guarded_licm_simt_0 +// The full invariant base chain must be hoisted in front of scf.for. +// CHECK: %[[TID:.*]] = pto.get_tid_x : i32 +// CHECK: %[[LANE:.*]] = arith.andi %[[TID]], %{{.*}} : i32 +// CHECK: %[[W64:.*]] = arith.extsi %{{.*}} : i32 to i64 +// CHECK: %[[SLOT:.*]] = arith.andi %[[W64]], %{{.*}} : i64 +// CHECK: %[[SLOTBASE:.*]] = arith.muli %[[SLOT]], %{{.*}} : i64 +// CHECK: %[[VALID64:.*]] = arith.extsi %{{.*}} : i32 to i64 +// CHECK: %[[LOCAL64:.*]] = arith.extsi %{{.*}} : i32 to i64 +// CHECK: %[[RANGEBASE:.*]] = arith.minsi %[[VALID64]], %[[LOCAL64]] : i64 +// CHECK: %[[TID64:.*]] = arith.extsi %[[TID]] : i32 to i64 +// CHECK: %[[LANE64:.*]] = arith.andi %[[TID64]], %{{.*}} : i64 +// CHECK: %[[BASE0:.*]] = arith.addi %[[SLOTBASE]], %[[RANGEBASE]] : i64 +// CHECK: %[[BASE1:.*]] = arith.addi %[[BASE0]], %[[LANE64]] : i64 +// CHECK: %[[BASE2:.*]] = arith.addi %[[BASE1]], %{{.*}} : i64 +// CHECK: %[[BASE:.*]] = arith.subi %[[BASE2]], %{{.*}} : i64 +// CHECK: scf.for +// The guard keeps only the IV-dependent i*32 portion, reusing the hoisted +// base and lane values. +// CHECK: %[[IV32:.*]] = arith.index_cast {{.*}} : index to i32 +// CHECK: %[[ITEROFF:.*]] = arith.muli %[[IV32]], %{{.*}} : i32 +// CHECK: %[[LOGICAL:.*]] = arith.addi %[[ITEROFF]], %[[LANE]] : i32 +// CHECK: arith.cmpi +// CHECK: scf.if +// CHECK: %[[IV64:.*]] = arith.extsi %[[IV32]] : i32 to i64 +// CHECK: %[[DYNAMIC:.*]] = arith.muli %[[IV64]], %{{.*}} : i64 +// CHECK: %[[OFFSET:.*]] = arith.addi %[[BASE]], %[[DYNAMIC]] : i64 +// CHECK: %[[INDEX:.*]] = arith.index_cast %[[OFFSET]] : i64 to index +// CHECK: pto.store %{{.*}}, %{{.*}}[%[[INDEX]]] diff --git a/test/lit/vpto/vpto_guarded_licm_pipeline.pto b/test/lit/vpto/vpto_guarded_licm_pipeline.pto new file mode 100644 index 0000000000..2575f190b8 --- /dev/null +++ b/test/lit/vpto/vpto_guarded_licm_pipeline.pto @@ -0,0 +1,65 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// End-to-end regression for the guarded-LICM gap: the pto.section.simt body is +// outlined into a helper, and inside the helper a loop-invariant base address +// chain is computed in an scf.if guarded by an IV-dependent condition. The +// full VPTO emission pipeline (which runs VPTOGuardedLICM between soft +// post-update and the generic LICM) must hoist the invariant chain in front +// of the scf.for while the IV-dependent offset chain and the guarded store +// stay inside the guard. +// +// RUN: ptoas --pto-backend=vpto --pto-arch=a5 --emit-vpto %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @guarded_licm_kernel(%dst: !pto.ptr, %laneid: i32) attributes {pto.entry} { + pto.section.simt<<<32, 1, 1>>> { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + %c1024 = arith.constant 1024 : index + %c128_i64 = arith.constant 128 : i64 + scf.for %iv = %c0 to %c1024 step %c64 { + %in_bounds = arith.cmpi slt, %iv, %c1024 : index + scf.if %in_bounds { + // Loop-invariant base chain (%laneid -> slot): hoisted to the loop + // preheader by VPTOGuardedLICM. + %w64 = arith.extsi %laneid : i32 to i64 + %slot = arith.muli %w64, %c128_i64 : i64 + // IV-dependent offset chain: must remain inside the guard. + %iv64 = arith.index_cast %iv : index to i64 + %dynamic = arith.muli %iv64, %c128_i64 : i64 + %offset = arith.addi %slot, %dynamic : i64 + %addr = arith.index_cast %offset : i64 to index + pto.store %laneid, %dst[%addr] : !pto.ptr, i32 + } + } + } + return + } +} + +// CHECK-LABEL: func.func @guarded_licm_kernel +// CHECK: pto.store_vfsimt_info +// CHECK: call @guarded_licm_kernel_simt_0 + +// CHECK-LABEL: func.func private @guarded_licm_kernel_simt_0 +// The invariant base chain must appear before scf.for... +// CHECK: %[[W64:.*]] = arith.extsi %{{.*}} : i32 to i64 +// CHECK: %[[SLOT:.*]] = arith.muli %[[W64]], %{{.*}} : i64 +// ...while no IV-dependent index computation may appear before the loop. +// CHECK-NOT: arith.index_cast +// CHECK: scf.for +// CHECK: arith.cmpi +// CHECK: scf.if +// The IV-dependent offset chain stays inside the guard, reusing the hoisted +// base chain in the final add. +// CHECK: %[[IV64:.*]] = arith.index_cast %{{.*}} : index to i64 +// CHECK: %[[DYN:.*]] = arith.muli %[[IV64]], %{{.*}} : i64 +// CHECK: %[[OFF:.*]] = arith.addi %[[SLOT]], %[[DYN]] : i64 +// CHECK: %[[ADDR:.*]] = arith.index_cast %[[OFF]] : i64 to index +// CHECK: pto.store %{{.*}}, %{{.*}}[%[[ADDR]]] diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index d57a82ee2b..fd54bdc6fe 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3122,6 +3122,10 @@ static void prepareVPTOForEmission(PassManager &pm) { if (enableSoftPostUpdate) { kernelModulePM.addPass(pto::createVPTOSoftPostUpdatePass()); } + // Hoist loop-invariant guarded address chains out of scf.if regions before + // the generic LICM (which only inspects top-level loop-body operations). + kernelModulePM.addNestedPass( + pto::createVPTOGuardedLICMPass()); kernelModulePM.addPass(createLoopInvariantCodeMotionPass()); kernelModulePM.addNestedPass( pto::createPTONarrowVPTOLoopCountersPass());