Expand a whole-aggregate store into per-field stores

src/enzyme_ad/jax/Passes/LLVMToAffineAccess.cpp

@@ -947,6 +947,204 @@ struct SplitAggregateLoad : public OpRewritePattern<affine::AffineLoadOp> {
   }
 };
 
+// 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<affine::AffineStoreOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  static void
+  collectLeaves(Type ty, SmallVectorImpl<int64_t> &path,
+                SmallVectorImpl<std::pair<SmallVector<int64_t>, Type>> &out) {
+    if (auto ST = dyn_cast<LLVM::LLVMStructType>(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<LLVM::LLVMArrayType>(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<int64_t>(path.begin(), path.end()), ty});
+  }
+
+  LogicalResult matchAndRewrite(affine::AffineStoreOp st,
+                                PatternRewriter &rewriter) const override {
+    auto structTy =
+        dyn_cast<LLVM::LLVMStructType>(st.getValueToStore().getType());
+    if (!structTy)
+      return failure();
+    auto p2m = st.getMemRef().getDefiningOp<enzymexla::Pointer2MemrefOp>();
+    if (!p2m)
+      return failure();
+    auto MT = cast<MemRefType>(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<std::pair<SmallVector<int64_t>, Type>> leaves;
+    SmallVector<int64_t> path;
+    collectLeaves(structTy, path, leaves);
+    if (leaves.empty())
+      return failure();
+
+    struct Plan {
+      SmallVector<int64_t> path;
+      uint64_t off;
+      Type ty;
+    };
+    SmallVector<Plan> 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<LLVM::AllocaOp> {
+  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<Access> accesses;
+    for (Operation *user : alloca->getUsers()) {
+      if (isa<LLVM::LifetimeStartOp, LLVM::LifetimeEndOp>(user))
+        continue;
+      auto p2m = dyn_cast<enzymexla::Pointer2MemrefOp>(user);
+      if (!p2m)
+        return failure();
+      auto MT = cast<MemRefType>(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<affine::AffineLoadOp>(vu)) {
+          map = ld.getAffineMap();
+          isStore = false;
+        } else if (auto st = dyn_cast<affine::AffineStoreOp>(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<AffineConstantExpr>(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<affine::AffineStoreOp>(only);
+      auto load = cast<affine::AffineLoadOp>(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<affine::AffineLoadOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -2280,8 +2478,8 @@ convertLLVMToAffineAccess(Operation *op,
         SimplifyDeadAlloc<memref::AllocaOp>, SimplifyDeadAlloc<memref::AllocOp>,
         SimplifyDeadAlloc<LLVM::AllocaOp>,
         SimplifyDeadAlloc<gpu::AllocOp, true>, Pointer2MemrefSelect, LoadSelect,
-        SplitAggregateLoad, AffineIfDeadResults,
-        SimpleMem2Reg<memref::AllocaOp>>(context);
+        SplitAggregateLoad, ExpandAggregateStore, ForwardSlotStores,
+        AffineIfDeadResults, SimpleMem2Reg<memref::AllocaOp>>(context);
     GreedyRewriteConfig config;
     config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Normal);
     config.enableFolding();

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<?xf64>) {
+    %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<?x!llvm.struct<(f64, i32, f64)>>
+    affine.store %upd, %view[0] : memref<?x!llvm.struct<(f64, i32, f64)>>
+    %fview = "enzymexla.pointer2memref"(%slot) : (!llvm.ptr) -> memref<?xf64>
+    %f0 = affine.load %fview[0] : memref<?xf64>
+    %f2 = affine.load %fview[2] : memref<?xf64>
+    %sum = arith.addf %f0, %f2 : f64
+    affine.store %sum, %out[0] : memref<?xf64>
+    return
+  }
+}
+
+// CHECK-LABEL: func.func @capture(
+// CHECK-SAME: %[[CAP:.+]]: !llvm.struct<(f64, i32, f64)>, %[[N:.+]]: i32, %[[OUT:.+]]: memref<?xf64>
+// 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<?xf64>
+// CHECK-NEXT: return
+// CHECK-NEXT: }