From dbb9f2f282d5e0250d593b4efa5169c2c2bdd637 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Tue, 25 Aug 2026 22:06:29 -0500 Subject: [PATCH] Expand a whole-aggregate store into per-field stores A device lambda that updates one capture field copies the whole capture into a stack slot through a single aggregate store of an insertvalue chain. Nothing decomposes an aggregate store (only aggregate loads), so the slot stays untyped and every kernel using it fails to raise. Expand the store into one typed store per leaf field through element-typed views; the extractvalues fold against the insertvalue chain, and the piece stores then forward to their loads. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD --- .../jax/Passes/LLVMToAffineAccess.cpp | 202 +++++++++++++++++- test/lit_tests/expand_aggregate_store.mlir | 30 +++ 2 files changed, 230 insertions(+), 2 deletions(-) create mode 100644 test/lit_tests/expand_aggregate_store.mlir diff --git a/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp b/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp index a024a977e0..b6083c13a3 100644 --- a/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp +++ b/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp @@ -947,6 +947,204 @@ struct SplitAggregateLoad : public OpRewritePattern { } }; +// A whole-aggregate store -- a lambda capture copied into a stack slot after +// an insertvalue update -- keeps the slot untyped and blocks every access +// analysis behind it. Expand it into one store per leaf field; the +// extractvalues fold against the insertvalue chain that built the value. +struct ExpandAggregateStore : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + static void + collectLeaves(Type ty, SmallVectorImpl &path, + SmallVectorImpl, Type>> &out) { + if (auto ST = dyn_cast(ty)) { + for (auto &&[i, member] : llvm::enumerate(ST.getBody())) { + path.push_back((int64_t)i); + collectLeaves(member, path, out); + path.pop_back(); + } + return; + } + if (auto AT = dyn_cast(ty)) { + for (uint64_t i = 0; i < AT.getNumElements(); ++i) { + path.push_back((int64_t)i); + collectLeaves(AT.getElementType(), path, out); + path.pop_back(); + } + return; + } + out.push_back({SmallVector(path.begin(), path.end()), ty}); + } + + LogicalResult matchAndRewrite(affine::AffineStoreOp st, + PatternRewriter &rewriter) const override { + auto structTy = + dyn_cast(st.getValueToStore().getType()); + if (!structTy) + return failure(); + auto p2m = st.getMemRef().getDefiningOp(); + if (!p2m) + return failure(); + auto MT = cast(p2m.getType()); + if (MT.getElementType() != structTy) + return failure(); + auto map = st.getAffineMap(); + if (map.getNumResults() != 1) + return failure(); + + DataLayout dl = DataLayout::closest(st); + uint64_t structSize = dl.getTypeSize(structTy); + + SmallVector, Type>> leaves; + SmallVector path; + collectLeaves(structTy, path, leaves); + if (leaves.empty()) + return failure(); + + struct Plan { + SmallVector path; + uint64_t off; + Type ty; + }; + SmallVector plans; + for (auto &[fpath, fty] : leaves) { + auto fo = fieldByteOffset(structTy, fpath, dl); + if (!fo) + return failure(); + auto [off, ty] = *fo; + uint64_t tsize = dl.getTypeSize(ty); + if (!tsize || off % tsize || structSize % tsize) + return failure(); + plans.push_back({fpath, off, ty}); + } + + rewriter.setInsertionPoint(st); + for (auto &plan : plans) { + uint64_t tsize = dl.getTypeSize(plan.ty); + Value v = LLVM::ExtractValueOp::create(rewriter, st.getLoc(), + st.getValueToStore(), plan.path); + AffineMap newMap = + AffineMap::get(map.getNumDims(), map.getNumSymbols(), + map.getResult(0) * (int64_t)(structSize / tsize) + + (int64_t)(plan.off / tsize)); + auto view = enzymexla::Pointer2MemrefOp::create( + rewriter, p2m.getLoc(), + MemRefType::get({ShapedType::kDynamic}, plan.ty, + MemRefLayoutAttrInterface{}, MT.getMemorySpace()), + p2m.getOperand()); + affine::AffineStoreOp::create(rewriter, st.getLoc(), v, view, newMap, + st.getMapOperands()); + } + rewriter.eraseOp(st); + return success(); + } +}; + +// The expanded piece stores land after the last mem2reg of the pipeline, so +// forward them here. A slot whose every access is a constant offset through +// a view is a bundle of registers: an offset written exactly once forwards +// to every load of the same offset and type the store dominates. +struct ForwardSlotStores : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(LLVM::AllocaOp alloca, + PatternRewriter &rewriter) const override { + DataLayout dl = DataLayout::closest(alloca); + struct Access { + Operation *op; + uint64_t off, size; + bool isStore; + }; + SmallVector accesses; + for (Operation *user : alloca->getUsers()) { + if (isa(user)) + continue; + auto p2m = dyn_cast(user); + if (!p2m) + return failure(); + auto MT = cast(p2m.getType()); + if (MT.getRank() != 1) + return failure(); + uint64_t esize = dl.getTypeSize(MT.getElementType()); + if (!esize) + return failure(); + for (Operation *vu : p2m->getUsers()) { + AffineMap map; + bool isStore; + if (auto ld = dyn_cast(vu)) { + map = ld.getAffineMap(); + isStore = false; + } else if (auto st = dyn_cast(vu)) { + if (st.getValueToStore() == p2m.getResult()) + return failure(); + map = st.getAffineMap(); + isStore = true; + } else { + return failure(); + } + if (map.getNumResults() != 1) + return failure(); + auto cst = dyn_cast(map.getResult(0)); + if (!cst || cst.getValue() < 0) + return failure(); + accesses.push_back( + {vu, (uint64_t)cst.getValue() * esize, esize, isStore}); + } + } + // Partially overlapping ranges make reaching values ambiguous. + for (auto &a : accesses) + for (auto &b : accesses) { + if (&a == &b) + continue; + bool disjoint = a.off + a.size <= b.off || b.off + b.size <= a.off; + bool identical = a.off == b.off && a.size == b.size; + if (!disjoint && !identical) + return failure(); + } + DominanceInfo DI(alloca->getParentOp()); + bool changed = false; + for (auto &ld : accesses) { + if (ld.isStore) + continue; + Operation *only = nullptr; + bool multiple = false; + for (auto &st : accesses) { + if (!st.isStore || st.off != ld.off) + continue; + if (only) + multiple = true; + only = st.op; + } + if (!only || multiple) + continue; + auto st = cast(only); + auto load = cast(ld.op); + if (st.getValueToStore().getType() != load.getType()) + continue; + if (!DI.properlyDominates(only, ld.op)) + continue; + rewriter.replaceOp(ld.op, st.getValueToStore()); + ld.op = nullptr; + changed = true; + } + // The slot cannot escape (every user was accounted for above), so a + // store no remaining load reads is dead. + for (auto &st : accesses) { + if (!st.isStore) + continue; + bool observed = false; + for (auto &ld : accesses) + if (!ld.isStore && ld.op && ld.off == st.off) + observed = true; + if (!observed) { + rewriter.eraseOp(st.op); + changed = true; + } + } + return success(changed); + } +}; + struct LoadSelect : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -2280,8 +2478,8 @@ convertLLVMToAffineAccess(Operation *op, SimplifyDeadAlloc, SimplifyDeadAlloc, SimplifyDeadAlloc, SimplifyDeadAlloc, Pointer2MemrefSelect, LoadSelect, - SplitAggregateLoad, AffineIfDeadResults, - SimpleMem2Reg>(context); + SplitAggregateLoad, ExpandAggregateStore, ForwardSlotStores, + AffineIfDeadResults, SimpleMem2Reg>(context); GreedyRewriteConfig config; config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Normal); config.enableFolding(); diff --git a/test/lit_tests/expand_aggregate_store.mlir b/test/lit_tests/expand_aggregate_store.mlir new file mode 100644 index 0000000000..a1528c7f8b --- /dev/null +++ b/test/lit_tests/expand_aggregate_store.mlir @@ -0,0 +1,30 @@ +// RUN: enzymexlamlir-opt %s --llvm-to-affine-access | FileCheck %s + +// A lambda capture copied whole into a stack slot after an insertvalue +// update: the aggregate store expands into per-field stores, the +// extractvalues fold against the insertvalue chain, the piece stores +// forward to their loads, and the dead slot disappears. +module { + func.func @capture(%cap: !llvm.struct<(f64, i32, f64)>, %n: i32, %out: memref) { + %c1 = arith.constant 1 : i32 + %upd = llvm.insertvalue %n, %cap[1] : !llvm.struct<(f64, i32, f64)> + %slot = llvm.alloca %c1 x !llvm.struct<(f64, i32, f64)> : (i32) -> !llvm.ptr + %view = "enzymexla.pointer2memref"(%slot) : (!llvm.ptr) -> memref> + affine.store %upd, %view[0] : memref> + %fview = "enzymexla.pointer2memref"(%slot) : (!llvm.ptr) -> memref + %f0 = affine.load %fview[0] : memref + %f2 = affine.load %fview[2] : memref + %sum = arith.addf %f0, %f2 : f64 + affine.store %sum, %out[0] : memref + return + } +} + +// CHECK-LABEL: func.func @capture( +// CHECK-SAME: %[[CAP:.+]]: !llvm.struct<(f64, i32, f64)>, %[[N:.+]]: i32, %[[OUT:.+]]: memref +// CHECK-NEXT: %[[F0:.+]] = llvm.extractvalue %[[CAP]][0] : !llvm.struct<(f64, i32, f64)> +// CHECK-NEXT: %[[F2:.+]] = llvm.extractvalue %[[CAP]][2] : !llvm.struct<(f64, i32, f64)> +// CHECK-NEXT: %[[SUM:.+]] = arith.addf %[[F0]], %[[F2]] : f64 +// CHECK-NEXT: affine.store %[[SUM]], %[[OUT]][0] : memref +// CHECK-NEXT: return +// CHECK-NEXT: }