From 530c251d51ee06c69b6f0e7b1cf7df6c2617f114 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Fri, 28 Aug 2026 09:50:27 -0500 Subject: [PATCH] Raising: buffer normalizations for viewed geps and raw accesses Kernels that are not fully inlined reach the raising with buffers addressed through pointer plumbing the tensor semantics cannot stand for. Normalize before raising: - rebaseViewedGeps: a typed view carved out of a buffer at an offset (gep feeding a pointer2memref) rebases its accesses onto the underlying buffer at base plus offset. - convertRawGepAccesses: data-dependent indexing (CSR-style loops over runtime offsets) can never become affine and stays as raw gep+load; the access still addresses whole elements, so it becomes a plain memref access through a flat view, which raising gathers. - inlineAllocaScopes/linearizeRegionBlocks: alloca scopes only delimit stack lifetime, meaningless under value semantics; splice them into the parent so scratch normalizations see the whole function. - dropDeadPointerChains: sweep the pointer plumbing the rewrites strand. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD --- .../jax/Passes/AffineToStableHLORaising.cpp | 339 +++++++++++++++++- test/lit_tests/raising/raw_gep_gather.mlir | 23 ++ test/lit_tests/raising/rebase_viewed_gep.mlir | 24 ++ 3 files changed, 384 insertions(+), 2 deletions(-) create mode 100644 test/lit_tests/raising/raw_gep_gather.mlir create mode 100644 test/lit_tests/raising/rebase_viewed_gep.mlir diff --git a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp index e3346bd36e..5410ef9496 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" @@ -3596,6 +3597,321 @@ 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. + // An alloca scope only delimits stack lifetime, which the raised value + // semantics make meaningless: splice its body into the parent. + // 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; + } + } + } + + // A barrier under a parallel axis of dynamic extent raises serialized, so + // it cannot be dropped as a no-op: distribute the loops around it first, + // cpuify-style, so every pre-barrier phase completes for the whole axis + // before the next phase starts. + + // A view taken of a gep result pins the kernel operand to the gep, which + // no tensor can stand for. Rebase the view onto the underlying pointer and + // fold the gep's element offset into each access index; data-dependent + // offsets make the accesses plain memref ops, which raising gathers. + static void rebaseViewedGeps(Operation *root) { + SmallVector geps; + root->walk([&](LLVM::GEPOp g) { geps.push_back(g); }); + for (auto gep : geps) { + auto idxs = gep.getIndices(); + if (idxs.size() != 1) + continue; + DataLayout dl = DataLayout::closest(gep); + int64_t elemSz = dl.getTypeSize(gep.getElemType()); + bool ok = true; + SmallVector views; + for (Operation *u : gep->getUsers()) { + auto p2m = dyn_cast(u); + if (!p2m || p2m.getType().getRank() != 1 || + !p2m.getType().getElementType().isIntOrFloat() || + (int64_t)dl.getTypeSize(p2m.getType().getElementType()) != elemSz) { + ok = false; + break; + } + for (Operation *a : p2m->getUsers()) { + if (auto ld = dyn_cast(a)) { + if (ld.getMap().getNumResults() == 1) + continue; + } else if (auto st = dyn_cast(a)) { + if (st.getMap().getNumResults() == 1 && + st.getValueToStore() != p2m.getResult()) + continue; + } else if (auto ld = dyn_cast(a)) { + if (ld.getIndices().size() == 1) + continue; + } else if (auto st = dyn_cast(a)) { + if (st.getIndices().size() == 1 && + st.getValueToStore() != p2m.getResult()) + continue; + } + ok = false; + break; + } + if (!ok) + break; + views.push_back(p2m); + } + if (!ok || views.empty()) + continue; + OpBuilder gb(gep); + Location loc = gep.getLoc(); + Value off; + if (!gep.getDynamicIndices().empty()) { + off = gep.getDynamicIndices()[0]; + if (!isa(off.getType())) + off = arith::IndexCastOp::create(gb, loc, gb.getIndexType(), off); + } else { + off = arith::ConstantIndexOp::create( + gb, loc, cast(idxs[0]).getInt()); + } + for (auto p2m : views) { + OpBuilder vb(p2m); + Value newView = enzymexla::Pointer2MemrefOp::create( + vb, p2m.getLoc(), p2m.getType(), gep.getBase()); + for (Operation *a : llvm::make_early_inc_range(p2m->getUsers())) { + OpBuilder ab(a); + auto toIdx = [&](AffineMap map, ValueRange operands) -> Value { + auto expanded = + affine::expandAffineMap(ab, a->getLoc(), map, operands); + return (*expanded)[0]; + }; + if (auto ld = dyn_cast(a)) { + Value idx = toIdx(ld.getMap(), ld.getMapOperands()); + idx = arith::AddIOp::create(ab, a->getLoc(), idx, off); + Value nl = memref::LoadOp::create(ab, a->getLoc(), newView, + ValueRange{idx}); + a->getResult(0).replaceAllUsesWith(nl); + a->erase(); + } else if (auto st = dyn_cast(a)) { + Value idx = toIdx(st.getMap(), st.getMapOperands()); + idx = arith::AddIOp::create(ab, a->getLoc(), idx, off); + memref::StoreOp::create(ab, a->getLoc(), st.getValueToStore(), + newView, ValueRange{idx}); + a->erase(); + } else if (auto ld = dyn_cast(a)) { + Value idx = + arith::AddIOp::create(ab, a->getLoc(), ld.getIndices()[0], off); + Value nl = memref::LoadOp::create(ab, a->getLoc(), newView, + ValueRange{idx}); + a->getResult(0).replaceAllUsesWith(nl); + a->erase(); + } else { + auto st = cast(a); + Value idx = + arith::AddIOp::create(ab, a->getLoc(), st.getIndices()[0], off); + memref::StoreOp::create(ab, a->getLoc(), st.getValueToStore(), + newView, ValueRange{idx}); + a->erase(); + } + } + p2m.erase(); + } + if (gep->use_empty()) + gep.erase(); + } + } + + // Data-dependent indexing (CSR-style loops over runtime offsets) can never + // become affine, so llvm-to-affine-access leaves it as raw gep+load. The + // access still addresses whole elements of the loaded type; a plain memref + // access through a flat view carries that, and raising gathers it. + static void convertRawGepAccesses(Operation *root) { + SmallVector accesses; + root->walk([&](Operation *op) { + if (isa(op)) + accesses.push_back(op); + }); + for (Operation *op : accesses) { + bool isLoad = isa(op); + if (isLoad ? cast(op).getVolatile_() + : cast(op).getVolatile_()) + continue; + Value addr = isLoad ? op->getOperand(0) : op->getOperand(1); + Type valTy = + isLoad ? op->getResult(0).getType() : op->getOperand(0).getType(); + if (!valTy.isIntOrFloat()) + continue; + DataLayout dl = DataLayout::closest(op); + Value base; + Value dynIdx; + int64_t constIdx = 0; + if (auto gep = addr.getDefiningOp()) { + auto idxs = gep.getIndices(); + if (idxs.size() != 1 || (int64_t)dl.getTypeSize(gep.getElemType()) != + (int64_t)dl.getTypeSize(valTy)) + continue; + base = gep.getBase(); + if (!gep.getDynamicIndices().empty()) + dynIdx = gep.getDynamicIndices()[0]; + else + constIdx = cast(idxs[0]).getInt(); + } else { + continue; + } + auto basePtrTy = cast(base.getType()); + Attribute space; + if (basePtrTy.getAddressSpace() != 0) + space = IntegerAttr::get(IntegerType::get(op->getContext(), 64), + basePtrTy.getAddressSpace()); + OpBuilder b(op); + Location loc = op->getLoc(); + auto MT = MemRefType::get({ShapedType::kDynamic}, valTy, + MemRefLayoutAttrInterface{}, space); + Value view = enzymexla::Pointer2MemrefOp::create(b, loc, MT, base); + Value idx; + if (dynIdx) { + idx = dynIdx; + if (!isa(idx.getType())) + idx = arith::IndexCastOp::create(b, loc, b.getIndexType(), idx); + } else { + idx = arith::ConstantIndexOp::create(b, loc, constIdx); + } + if (isLoad) { + Value ld = memref::LoadOp::create(b, loc, view, ValueRange{idx}); + op->getResult(0).replaceAllUsesWith(ld); + op->erase(); + } else { + memref::StoreOp::create(b, loc, op->getOperand(0), view, + ValueRange{idx}); + op->erase(); + } + } + } + + // Access rewrites leave dead pointer plumbing behind, and raising visits + // every op in the region: sweep the unused chains. + static void dropDeadPointerChains(Operation *root) { + bool changed = true; + while (changed) { + changed = false; + SmallVector dead; + root->walk([&](Operation *op) { + if (isa(op) && + op->use_empty()) + dead.push_back(op); + }); + for (Operation *op : dead) { + op->erase(); + changed = true; + } + } + } + static void stripAccessMemorySpaceCasts(Operation *root) { SmallVector casts; root->walk([&](memref::MemorySpaceCastOp c) { casts.push_back(c); }); @@ -3899,7 +4215,13 @@ struct AffineToStableHLORaisingPass // Peeling rewrites loops, so it stays scoped to the regions this pass // actually raises. for (auto func : funcs) { - stripAccessMemorySpaceCasts(func); + inlineAllocaScopes(func); + for (int round = 0; round < 2; ++round) { + stripAccessMemorySpaceCasts(func); + rebaseViewedGeps(func); + convertRawGepAccesses(func); + } + dropDeadPointerChains(func); boundParallelAxes(func); peelDynamicParallelDims(func); } @@ -3922,7 +4244,20 @@ struct AffineToStableHLORaisingPass std::vector gwrap; op->walk([&](enzymexla::GPUWrapperOp g) { gwrap.push_back(g); }); for (auto g : gwrap) { - stripAccessMemorySpaceCasts(g); + // Scope inlining hoists scratch allocas to the surrounding function, + // so the buffer normalizations must see the whole function, not just + // the wrapper region; the rewrites also expose one another (a rebase + // creates the direct views a flatten wants), so iterate once more. + Operation *root = g->getParentOfType(); + if (!root) + root = g; + inlineAllocaScopes(root); + for (int round = 0; round < 2; ++round) { + stripAccessMemorySpaceCasts(root); + rebaseViewedGeps(root); + convertRawGepAccesses(root); + } + dropDeadPointerChains(root); boundParallelAxes(g); peelDynamicParallelDims(g); } diff --git a/test/lit_tests/raising/raw_gep_gather.mlir b/test/lit_tests/raising/raw_gep_gather.mlir new file mode 100644 index 0000000000..cacda9503a --- /dev/null +++ b/test/lit_tests/raising/raw_gep_gather.mlir @@ -0,0 +1,23 @@ +// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo | FileCheck %s + +// Data-dependent indexing (CSR-style) stays as raw gep+load; the access +// still addresses whole elements, so it converts to a flat memref access +// and raises as a gather. + +// CHECK-LABEL: @csr_raised +// CHECK: stablehlo.gather +// CHECK-NOT: llvm.load + +module { + func.func private @csr(%out: memref<16xf64, 1>, %in: memref<16xf64, 1>, %idx: memref<16xi32, 1>) { + affine.parallel (%t) = (0) to (16) { + %i = affine.load %idx[%t] : memref<16xi32, 1> + %i64 = arith.extsi %i : i32 to i64 + %p = "enzymexla.memref2pointer"(%in) : (memref<16xf64, 1>) -> !llvm.ptr<1> + %g = llvm.getelementptr %p[%i64] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, f64 + %v = llvm.load %g : !llvm.ptr<1> -> f64 + affine.store %v, %out[%t] : memref<16xf64, 1> + } + return + } +} diff --git a/test/lit_tests/raising/rebase_viewed_gep.mlir b/test/lit_tests/raising/rebase_viewed_gep.mlir new file mode 100644 index 0000000000..b7652d6960 --- /dev/null +++ b/test/lit_tests/raising/rebase_viewed_gep.mlir @@ -0,0 +1,24 @@ +// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo | FileCheck %s + +// A view carved out of a buffer at a runtime element offset (gep feeding a +// typed pointer2memref) rebases onto the underlying buffer: the view's +// accesses index the base plus the offset, and raise as a gather. + +// CHECK-LABEL: @viewed_raised +// CHECK: stablehlo.gather +// CHECK-NOT: llvm.getelementptr + +module { + func.func private @viewed(%out: memref<16xf64, 1>, %in: memref<64xf64, 1>, %idx: memref<16xi32, 1>) { + affine.parallel (%t) = (0) to (16) { + %i = affine.load %idx[%t] : memref<16xi32, 1> + %i64 = arith.extsi %i : i32 to i64 + %p = "enzymexla.memref2pointer"(%in) : (memref<64xf64, 1>) -> !llvm.ptr<1> + %g = llvm.getelementptr %p[%i64] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, f64 + %view = "enzymexla.pointer2memref"(%g) : (!llvm.ptr<1>) -> memref + %v = affine.load %view[0] : memref + affine.store %v, %out[%t] : memref<16xf64, 1> + } + return + } +}