diff --git a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp index e3346bd36e..20f3937445 100644 --- a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp +++ b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp @@ -22,6 +22,7 @@ #include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/Math/IR/Math.h" @@ -3415,6 +3416,17 @@ tryRaisingOpToStableHLO(Operation *op, IRMapping &mapping, OpBuilder &builder, return success(); } + // A shape-erasing cast of an already-raised buffer is the identity on the + // underlying tensor. + if (auto castOp = dyn_cast(op)) { + Value src = mapping.lookupOrNull(castOp.getSource()); + if (src && maps.count(src)) { + mapping.map(castOp.getResult(), src); + maps[src] = maps.lookup(src); + return success(); + } + } + // Raised execution is ordered over whole tensors: a store over a batched // thread axis completes for the entire axis before the next op runs, which // is exactly what the barrier guaranteed. @@ -3596,6 +3608,121 @@ struct AffineToStableHLORaisingPass // raising identifies buffers by SSA root: a memory_space_cast view would // split one buffer into two roots and lose store propagation. Retarget the // accesses to the source and drop the cast. + // Straight-line CFG inside a cloned callee region folds into one block, + // so the scope inlining below can dissolve it. + static void linearizeRegionBlocks(Region &r) { + auto isTrapBlock = [](Block *b) { + return isa(b->getTerminator()); + }; + bool changed = true; + while (changed) { + changed = false; + // Branches whose other targets only trap take their one live successor + // unconditionally. + for (Block &b : r) { + Operation *term = b.getTerminator(); + SmallVector>> live; + bool anyTrap = false; + if (isa(term)) { + for (auto [i, succ] : llvm::enumerate(term->getSuccessors())) { + if (isTrapBlock(succ)) { + anyTrap = true; + continue; + } + auto sops = cast(term).getSuccessorOperands(i); + SmallVector args(sops.getForwardedOperands().begin(), + sops.getForwardedOperands().end()); + live.push_back({succ, std::move(args)}); + } + } + if (anyTrap && live.size() == 1) { + OpBuilder tb(term); + cf::BranchOp::create(tb, term->getLoc(), live[0].first, + live[0].second); + term->erase(); + changed = true; + break; + } + } + if (changed) + continue; + // Trap blocks with no remaining predecessors disappear. + for (Block &b : llvm::make_early_inc_range(r)) { + if (&b != &r.front() && b.hasNoPredecessors()) { + b.dropAllDefinedValueUses(); + b.erase(); + changed = true; + } + } + if (changed) + continue; + for (Block &b : r) { + Operation *term = b.getTerminator(); + Block *succ = nullptr; + SmallVector args; + if (auto br = dyn_cast(term)) { + succ = br.getDest(); + args.assign(br.getDestOperands().begin(), br.getDestOperands().end()); + } else if (auto br = dyn_cast(term)) { + succ = br.getDest(); + args.assign(br.getDestOperands().begin(), br.getDestOperands().end()); + } else { + continue; + } + if (!succ || succ == &b || succ->getSinglePredecessor() != &b) + continue; + for (auto [ba, v] : llvm::zip(succ->getArguments(), args)) + ba.replaceAllUsesWith(v); + term->erase(); + b.getOperations().splice(b.end(), succ->getOperations()); + succ->erase(); + changed = true; + break; + } + } + } + + static void inlineAllocaScopes(Operation *g) { + // Inliner wrappers stack alloca_scope/execute_region pairs, so inlining + // one can expose another: iterate to a fixed point. + bool changed = true; + while (changed) { + changed = false; + SmallVector scopes; + g->walk([&](Operation *op) { + if (isa(op)) + scopes.push_back(op); + }); + for (Operation *sc : scopes) { + Region &r = sc->getRegion(0); + if (!r.hasOneBlock()) + linearizeRegionBlocks(r); + if (!r.hasOneBlock()) { + if (getenv("DEBUG_SCOPES")) { + llvm::errs() << "scope multiblock after linearize: " + << std::distance(r.begin(), r.end()) << " blocks;"; + for (Block &b : r) + llvm::errs() << " term=" << b.getTerminator()->getName() + << " preds=" + << std::distance(b.pred_begin(), b.pred_end()); + llvm::errs() << "\n"; + } + continue; + } + Block *body = &r.front(); + Operation *term = body->getTerminator(); + for (auto [res, yielded] : + llvm::zip(sc->getResults(), term->getOperands())) + res.replaceAllUsesWith(yielded); + term->erase(); + sc->getBlock()->getOperations().splice(sc->getIterator(), + body->getOperations()); + sc->erase(); + changed = true; + } + } + } + static void stripAccessMemorySpaceCasts(Operation *root) { SmallVector casts; root->walk([&](memref::MemorySpaceCastOp c) { casts.push_back(c); }); @@ -3609,6 +3736,25 @@ struct AffineToStableHLORaisingPass u->replaceUsesOfWith(c.getResult(), c.getSource()); c.erase(); } + // Shape-erasing casts of static buffers block raising the same way: + // accesses go straight to the static source. + SmallVector shapeCasts; + root->walk([&](memref::CastOp c) { + auto src = dyn_cast(c.getSource().getType()); + auto dst = dyn_cast(c.getType()); + if (src && dst && src.hasStaticShape() && !dst.hasStaticShape()) + shapeCasts.push_back(c); + }); + for (auto c : shapeCasts) { + if (!llvm::all_of(c->getUsers(), [](Operation *u) { + return isa(u); + })) + continue; + for (Operation *u : llvm::make_early_inc_range(c->getUsers())) + u->replaceUsesOfWith(c.getResult(), c.getSource()); + c.erase(); + } } // A parallel dimension whose extent is only known at runtime cannot become @@ -3900,6 +4046,7 @@ struct AffineToStableHLORaisingPass // actually raises. for (auto func : funcs) { stripAccessMemorySpaceCasts(func); + inlineAllocaScopes(func); boundParallelAxes(func); peelDynamicParallelDims(func); } @@ -3923,6 +4070,7 @@ struct AffineToStableHLORaisingPass op->walk([&](enzymexla::GPUWrapperOp g) { gwrap.push_back(g); }); for (auto g : gwrap) { stripAccessMemorySpaceCasts(g); + inlineAllocaScopes(g); boundParallelAxes(g); peelDynamicParallelDims(g); } diff --git a/test/lit_tests/raising/raise_scope_cast.mlir b/test/lit_tests/raising/raise_scope_cast.mlir new file mode 100644 index 0000000000..f5d5232656 --- /dev/null +++ b/test/lit_tests/raising/raise_scope_cast.mlir @@ -0,0 +1,47 @@ +// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo --split-input-file | FileCheck %s + +// A shape-erasing cast of static scratch only renames the buffer; accesses +// go straight to the static source. +func.func @castuse(%out: memref<16xf64, 1>, %in: memref<16xf64, 1>) { + %scr = memref.alloca() : memref<16xf64> + %dyn = memref.cast %scr : memref<16xf64> to memref + affine.parallel (%t) = (0) to (16) { + %v = affine.load %in[%t] : memref<16xf64, 1> + affine.store %v, %dyn[%t] : memref + %r = affine.load %dyn[15 - %t] : memref + affine.store %r, %out[%t] : memref<16xf64, 1> + } + return +} + +// CHECK-LABEL: func.func private @castuse_raised( +// CHECK-NOT: memref.cast +// CHECK: stablehlo.reverse + +// ----- + +// An inliner-wrapped callee arrives as an execute_region whose CFG carries a +// trap arm; the branch takes its one live successor, the straight line +// merges, and the scope dissolves into the surrounding kernel. +func.func @scoped(%out: memref<16xi32, 1>, %in: memref<16xi32, 1>) { + affine.parallel (%t) = (0) to (16) { + %v = affine.load %in[%t] : memref<16xi32, 1> + %r = scf.execute_region -> i32 { + %c100 = arith.constant 100 : i32 + %c = arith.cmpi slt, %v, %c100 : i32 + cf.cond_br %c, ^bb1, ^bb2 + ^bb1: + %c1 = arith.constant 1 : i32 + %s = arith.addi %v, %c1 : i32 + scf.yield %s : i32 + ^bb2: + llvm.unreachable + } + affine.store %r, %out[%t] : memref<16xi32, 1> + } + return +} + +// CHECK-LABEL: func.func private @scoped_raised( +// CHECK-NOT: scf.execute_region +// CHECK: stablehlo.add