diff --git a/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp b/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp index e84445f8f0..f4f74f8df9 100644 --- a/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp +++ b/src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp @@ -949,6 +949,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; @@ -2521,8 +2719,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: }