From cbec48aff244e86e72aa69121a4fdac8d4128fc4 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Wed, 26 Aug 2026 11:12:37 -0500 Subject: [PATCH] Split struct-element kernel scratch into per-field buffers Reduce-with-location kernels (MFEM's ArgMin/ArgMax/MinMaxLoc reducers) keep their block partials in shared scratch of struct type -- memref<256 x struct<(f64, i32)>> -- which cannot become a tensor whole. Every access reaches the scratch through a flat primitive view whose index pins a fixed byte offset within the struct: affine maps with a constant residue, plain arith index chains with a derivable residue, whole-struct integer moves, and fixed-size memcpys between struct-strided geps. Split the array-of-structs into one primitive scratch per field and rewrite each access form onto its field buffer, with same-size type punning kept as bitcasts and struct-wide integer moves recomposed with shifts. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD --- fc.err | 30 ++ .../jax/Passes/AffineToStableHLORaising.cpp | 445 ++++++++++++++++++ .../raising/raise_struct_scratch.mlir | 67 +++ 3 files changed, 542 insertions(+) create mode 100644 fc.err create mode 100644 test/lit_tests/raising/raise_struct_scratch.mlir diff --git a/fc.err b/fc.err new file mode 100644 index 0000000000..ea5c2320c7 --- /dev/null +++ b/fc.err @@ -0,0 +1,30 @@ +test/lit_tests/raising/raise_struct_scratch.mlir:54:7: error: masked affine.store is dependent on less dimensions than masked stored value: +"affine.store"(%7, %1) <{map = affine_map<() -> (0)>}> : (f64, memref<256xf64>) -> () +: () -> (0) +%arg2, : (d0) -> (d0) + "llvm.intr.memcpy"(%dst, %src, %c16_i64) <{isVolatile = false}> : (!llvm.ptr, !llvm.ptr, i64) -> () + ^ +test/lit_tests/raising/raise_struct_scratch.mlir:54:7: note: see current operation: "affine.store"(%7, %1) <{map = affine_map<() -> (0)>}> : (f64, memref<256xf64>) -> () +failed to raise func: func.func @mixedpair(%arg0: memref<256xf64, 1>, %arg1: memref<256xf64, 1>) { + %c16_i64 = arith.constant 16 : i64 + %alloca = memref.alloca() : memref<256xf64> + %alloca_0 = memref.alloca() : memref<256xi32> + affine.parallel (%arg2) = (0) to (256) { + %0 = affine.load %arg1[%arg2] : memref<256xf64, 1> + affine.store %0, %alloca[%arg2] : memref<256xf64> + %1 = arith.index_castui %arg2 : index to i32 + affine.store %1, %alloca_0[%arg2] : memref<256xi32> + "enzymexla.barrier"(%arg2) : (index) -> () + affine.if affine_set<(d0) : (d0 == 0)>(%arg2) { + %c1 = arith.constant 1 : index + %3 = affine.load %alloca[1] : memref<256xf64> + affine.store %3, %alloca[0] : memref<256xf64> + %4 = affine.load %alloca_0[1] : memref<256xi32> + affine.store %4, %alloca_0[0] : memref<256xi32> + } + "enzymexla.barrier"(%arg2) : (index) -> () + %2 = affine.load %alloca[%arg2] : memref<256xf64> + affine.store %2, %arg0[%arg2] : memref<256xf64, 1> + } + return +} diff --git a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp index ece5725dce..2c17d11c01 100644 --- a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp +++ b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp @@ -3697,6 +3697,449 @@ struct AffineToStableHLORaisingPass } } + // Residue of an integer/index value modulo m, when the defining arith + // chain pins it down. Extension and truncation preserve residues for the + // power-of-two moduli struct layouts produce. + static std::optional staticResidue(Value v, int64_t m, + unsigned depth = 0) { + if (m == 1) + return 0; + if (depth > 16) + return std::nullopt; + APInt cst; + if (matchPattern(v, m_ConstantInt(&cst))) + return ((cst.getSExtValue() % m) + m) % m; + Operation *op = v.getDefiningOp(); + if (!op) + return std::nullopt; + if (isa(op) || + (isa(op) && llvm::isPowerOf2_64(m))) + return staticResidue(op->getOperand(0), m, depth + 1); + if (auto add = dyn_cast(op)) { + auto a = staticResidue(add.getLhs(), m, depth + 1); + auto b = staticResidue(add.getRhs(), m, depth + 1); + if (a && b) + return (*a + *b) % m; + return std::nullopt; + } + if (auto mul = dyn_cast(op)) { + auto a = staticResidue(mul.getLhs(), m, depth + 1); + auto b = staticResidue(mul.getRhs(), m, depth + 1); + if (a && b) + return (*a * *b) % m; + if ((a && *a == 0) || (b && *b == 0)) + return 0; + return std::nullopt; + } + if (auto shl = dyn_cast(op)) { + APInt sh; + if (matchPattern(shl.getRhs(), m_ConstantInt(&sh)) && + sh.getZExtValue() < 63) { + int64_t f = (int64_t(1) << sh.getZExtValue()) % m; + if (f == 0) + return 0; + auto a = staticResidue(shl.getLhs(), m, depth + 1); + if (a) + return (*a * f) % m; + } + return std::nullopt; + } + return std::nullopt; + } + + // Struct-element scratch (e.g. reduce-with-location value/index pairs) + // cannot become a tensor whole. Every access reaches it through a flat + // primitive view whose affine index resolves to a fixed byte offset within + // the struct, so the array-of-structs splits into one primitive scratch + // per field, with whole-struct integer moves split into their fields. + static void splitStructScratch(Operation *root) { + SmallVector allocas; + root->walk([&](memref::AllocaOp a) { allocas.push_back(a); }); + for (auto alloca : allocas) { + auto MT = alloca.getType(); + if (!MT.hasStaticShape()) + continue; + auto ST = dyn_cast(MT.getElementType()); + if (!ST) + continue; + DataLayout dl = DataLayout::closest(alloca); + struct Field { + uint64_t off, size; + Type ty; + }; + SmallVector fields; + uint64_t byte = 0; + bool ok = true; + for (Type member : ST.getBody()) { + if (!member.isIntOrFloat()) { + ok = false; + break; + } + if (!ST.isPacked()) + byte = llvm::alignTo(byte, dl.getTypeABIAlignment(member)); + fields.push_back({byte, dl.getTypeSize(member), member}); + byte += dl.getTypeSize(member); + } + if (!ok) + continue; + int64_t pairSize = dl.getTypeSize(ST); + + struct FieldAccess { + Operation *op; + unsigned field; + AffineMap pairMap; + bool bitcast; + }; + struct WideAccess { + Operation *op; + SmallVector covered; + uint64_t base; + AffineMap pairMap; + }; + struct DynAccess { + Operation *op; + unsigned field; + int64_t q; + bool bitcast; + }; + SmallVector fieldAccesses; + SmallVector wideAccesses; + SmallVector dynAccesses; + SmallVector views; + SmallVector casts; + + auto classify = [&](Operation *op, Type elemTy, AffineMap map) -> bool { + if (map.getNumResults() != 1) + return false; + int64_t s = dl.getTypeSize(elemTy); + AffineExpr byteExpr = map.getResult(0) * s; + AffineExpr rem = simplifyAffineExpr( + byteExpr % pairSize, map.getNumDims(), map.getNumSymbols()); + auto remCst = dyn_cast(rem); + if (!remCst) + return false; + uint64_t off = remCst.getValue(); + AffineMap pairMap = AffineMap::get( + map.getNumDims(), map.getNumSymbols(), + simplifyAffineExpr(byteExpr.floorDiv(pairSize), map.getNumDims(), + map.getNumSymbols())); + for (auto &&[i, f] : llvm::enumerate(fields)) + if (f.off == off && f.size == (uint64_t)s) { + fieldAccesses.push_back({op, (unsigned)i, pairMap, + /*bitcast=*/f.ty != elemTy}); + return true; + } + // A wider integer move covering whole fields splits into them. + if (!isa(elemTy)) + return false; + SmallVector covered; + for (auto &&[i, f] : llvm::enumerate(fields)) { + if (f.off + f.size <= off || f.off >= off + s) + continue; + if (f.off < off || f.off + f.size > off + s || + !isa(f.ty)) + return false; + covered.push_back(i); + } + if (covered.empty()) + return false; + wideAccesses.push_back({op, covered, off, pairMap}); + return true; + }; + + // The same classification for accesses whose index arrives through + // plain arithmetic instead of an affine map: the field is fixed when + // the index has a static residue modulo the per-struct element count. + auto classifyDyn = [&](Operation *op, Type elemTy, Value idx) -> bool { + int64_t s = dl.getTypeSize(elemTy); + if (!s || pairSize % s) + return false; + int64_t q = pairSize / s; + auto r = staticResidue(idx, q); + if (!r) + return false; + uint64_t off = (uint64_t)*r * s; + for (auto &&[i, f] : llvm::enumerate(fields)) + if (f.off == off && f.size == (uint64_t)s) { + dynAccesses.push_back({op, (unsigned)i, q, + /*bitcast=*/f.ty != elemTy}); + return true; + } + return false; + }; + + SmallVector spaceCasts; + llvm::MapVector pairGeps; + llvm::SetVector pairCopies; + DenseSet basePtrs; + bool viewedOnly = true; + for (Operation *user : alloca->getUsers()) { + auto m2p = dyn_cast(user); + if (!m2p) { + viewedOnly = false; + break; + } + casts.push_back(m2p); + basePtrs.insert(m2p.getResult()); + SmallVector viewUsers(m2p->getUsers()); + for (unsigned vi = 0; vi < viewUsers.size() && viewedOnly; ++vi) { + Operation *viewUser = viewUsers[vi]; + if (isa(viewUser)) { + spaceCasts.push_back(viewUser); + basePtrs.insert(viewUser->getResult(0)); + viewUsers.append(viewUser->getUsers().begin(), + viewUser->getUsers().end()); + continue; + } + // A copy of pair zero addresses the scratch base directly. + if (auto mc = dyn_cast(viewUser)) { + APInt len; + if (mc.getIsVolatile() || + !matchPattern(mc.getLen(), m_ConstantInt(&len)) || + len.getSExtValue() != pairSize) { + viewedOnly = false; + break; + } + pairCopies.insert(mc); + continue; + } + // A whole-struct move arrives as a fixed-size memcpy between + // struct-strided geps into the scratch: split it per field. + if (auto gep = dyn_cast(viewUser)) { + auto idxs = gep.getIndices(); + if (idxs.size() != 1 || + (int64_t)dl.getTypeSize(gep.getElemType()) != pairSize) { + viewedOnly = false; + break; + } + bool copiesOnly = true; + for (Operation *gu : gep->getUsers()) { + auto mc = dyn_cast(gu); + APInt len; + if (!mc || mc.getIsVolatile() || + !matchPattern(mc.getLen(), m_ConstantInt(&len)) || + len.getSExtValue() != pairSize) { + copiesOnly = false; + break; + } + pairCopies.insert(mc); + } + if (!copiesOnly) { + viewedOnly = false; + break; + } + Value gepIdx; + if (!gep.getDynamicIndices().empty()) { + gepIdx = gep.getDynamicIndices()[0]; + } else { + OpBuilder gb(gep); + gepIdx = arith::ConstantIndexOp::create( + gb, gep.getLoc(), cast(idxs[0]).getInt()); + } + pairGeps.insert({gep, gepIdx}); + continue; + } + auto p2m = dyn_cast(viewUser); + if (!p2m || p2m.getType().getRank() != 1 || + !p2m.getType().getElementType().isIntOrFloat()) { + viewedOnly = false; + break; + } + for (Operation *access : p2m->getUsers()) { + if (auto ld = dyn_cast(access)) { + if (classify(ld, ld.getType(), ld.getMap())) + continue; + } else if (auto st = dyn_cast(access)) { + if (st.getValueToStore() != p2m.getResult() && + classify(st, st.getValueToStore().getType(), st.getMap())) + continue; + } else if (auto mld = dyn_cast(access)) { + if (mld.getIndices().size() == 1 && + classifyDyn(mld, mld.getType(), mld.getIndices()[0])) + continue; + } else if (auto mst = dyn_cast(access)) { + if (mst.getMemRef() == p2m.getResult() && + mst.getValueToStore() != p2m.getResult() && + mst.getIndices().size() == 1 && + classifyDyn(mst, mst.getValueToStore().getType(), + mst.getIndices()[0])) + continue; + } + viewedOnly = false; + break; + } + if (!viewedOnly) + break; + views.push_back(p2m); + } + if (!viewedOnly) + break; + } + // Every pair copy must connect two classified geps or scratch bases. + auto copyEnd = [&](Value p) { + return basePtrs.contains(p) || + (p.getDefiningOp() && pairGeps.count(p.getDefiningOp())); + }; + for (Operation *mc : pairCopies) { + auto cp = cast(mc); + if (!copyEnd(cp.getDst()) || !copyEnd(cp.getSrc())) + viewedOnly = false; + } + if (!viewedOnly || (fieldAccesses.empty() && wideAccesses.empty())) + continue; + + OpBuilder b(alloca); + SmallVector fieldBufs; + for (auto &f : fields) + fieldBufs.push_back(memref::AllocaOp::create( + b, alloca.getLoc(), MemRefType::get({MT.getNumElements()}, f.ty))); + + for (auto &fa : fieldAccesses) { + if (auto ld = dyn_cast(fa.op)) { + OpBuilder ab(ld); + Value newLd = + affine::AffineLoadOp::create(ab, ld.getLoc(), fieldBufs[fa.field], + fa.pairMap, ld.getMapOperands()); + if (fa.bitcast) + newLd = + arith::BitcastOp::create(ab, ld.getLoc(), ld.getType(), newLd); + ld.getResult().replaceAllUsesWith(newLd); + ld.erase(); + } else { + auto st = cast(fa.op); + OpBuilder ab(st); + Value val = st.getValueToStore(); + if (fa.bitcast) + val = arith::BitcastOp::create(ab, st.getLoc(), fields[fa.field].ty, + val); + affine::AffineStoreOp::create(ab, st.getLoc(), val, + fieldBufs[fa.field], fa.pairMap, + st.getMapOperands()); + st.erase(); + } + } + for (auto &wa : wideAccesses) { + if (auto ld = dyn_cast(wa.op)) { + OpBuilder ab(ld); + Location loc = ld.getLoc(); + Type wideTy = ld.getType(); + Value acc = arith::ConstantOp::create(ab, loc, wideTy, + ab.getIntegerAttr(wideTy, 0)); + for (unsigned i : wa.covered) { + Value v = affine::AffineLoadOp::create( + ab, loc, fieldBufs[i], wa.pairMap, ld.getMapOperands()); + Value z = arith::ExtUIOp::create(ab, loc, wideTy, v); + uint64_t sh = (fields[i].off - wa.base) * 8; + if (sh) { + Value shv = arith::ConstantOp::create( + ab, loc, wideTy, ab.getIntegerAttr(wideTy, sh)); + z = arith::ShLIOp::create(ab, loc, z, shv); + } + acc = arith::OrIOp::create(ab, loc, acc, z); + } + ld.getResult().replaceAllUsesWith(acc); + ld.erase(); + } else { + auto st = cast(wa.op); + OpBuilder ab(st); + Location loc = st.getLoc(); + Value val = st.getValueToStore(); + Type wideTy = val.getType(); + for (unsigned i : wa.covered) { + Value part = val; + uint64_t sh = (fields[i].off - wa.base) * 8; + if (sh) { + Value shv = arith::ConstantOp::create( + ab, loc, wideTy, ab.getIntegerAttr(wideTy, sh)); + part = arith::ShRUIOp::create(ab, loc, part, shv); + } + part = arith::TruncIOp::create(ab, loc, fields[i].ty, part); + affine::AffineStoreOp::create(ab, loc, part, fieldBufs[i], + wa.pairMap, st.getMapOperands()); + } + st.erase(); + } + } + for (auto &da : dynAccesses) { + OpBuilder ab(da.op); + Location loc = da.op->getLoc(); + Value idx = isa(da.op) + ? cast(da.op).getIndices()[0] + : cast(da.op).getIndices()[0]; + Value pairIdx = idx; + if (da.q != 1) { + Value qc = arith::ConstantIndexOp::create(ab, loc, da.q); + pairIdx = arith::DivUIOp::create(ab, loc, idx, qc); + } + if (auto mld = dyn_cast(da.op)) { + Value newLd = memref::LoadOp::create(ab, loc, fieldBufs[da.field], + ValueRange{pairIdx}); + if (da.bitcast) + newLd = arith::BitcastOp::create(ab, loc, mld.getType(), newLd); + mld.getResult().replaceAllUsesWith(newLd); + mld.erase(); + } else { + auto mst = cast(da.op); + Value val = mst.getValueToStore(); + if (da.bitcast) + val = arith::BitcastOp::create(ab, loc, fields[da.field].ty, val); + memref::StoreOp::create(ab, loc, val, fieldBufs[da.field], + ValueRange{pairIdx}); + mst.erase(); + } + } + for (Operation *mc : pairCopies) { + auto cp = cast(mc); + OpBuilder ab(cp); + Location loc = cp.getLoc(); + // Constant pair indices stay affine so the accesses raise directly. + auto toIdx = [&](Value p) -> std::pair> { + if (basePtrs.contains(p)) + return {Value(), 0}; + Value v = pairGeps.find(p.getDefiningOp())->second; + APInt c; + if (matchPattern(v, m_ConstantInt(&c))) + return {Value(), c.getSExtValue()}; + if (!isa(v.getType())) + v = arith::IndexCastUIOp::create(ab, loc, ab.getIndexType(), v); + return {v, std::nullopt}; + }; + auto [dstIdx, dstC] = toIdx(cp.getDst()); + auto [srcIdx, srcC] = toIdx(cp.getSrc()); + for (auto &&[i, f] : llvm::enumerate(fields)) { + Value v; + if (srcC) + v = affine::AffineLoadOp::create( + ab, loc, fieldBufs[i], + AffineMap::getConstantMap(*srcC, ab.getContext()), + ValueRange()); + else + v = memref::LoadOp::create(ab, loc, fieldBufs[i], + ValueRange{srcIdx}); + if (dstC) + affine::AffineStoreOp::create( + ab, loc, v, fieldBufs[i], + AffineMap::getConstantMap(*dstC, ab.getContext()), + ValueRange()); + else + memref::StoreOp::create(ab, loc, v, fieldBufs[i], + ValueRange{dstIdx}); + } + cp.erase(); + } + for (auto &g : pairGeps) + g.first->erase(); + for (auto p2m : views) + p2m.erase(); + for (Operation *sc : llvm::reverse(spaceCasts)) + sc->erase(); + for (auto m2p : casts) + m2p.erase(); + alloca.erase(); + } + } + void runOnOperation() override { ParallelContext::Options options{enable_lockstep_for, dump_failed_lockstep, prefer_while_raising, @@ -3735,6 +4178,7 @@ struct AffineToStableHLORaisingPass // actually raises. for (auto func : funcs) { stripAccessMemorySpaceCasts(func); + splitStructScratch(func); peelDynamicParallelDims(func); } @@ -3757,6 +4201,7 @@ struct AffineToStableHLORaisingPass op->walk([&](enzymexla::GPUWrapperOp g) { gwrap.push_back(g); }); for (auto g : gwrap) { stripAccessMemorySpaceCasts(g); + splitStructScratch(g); peelDynamicParallelDims(g); } size_t raised_count = 0; diff --git a/test/lit_tests/raising/raise_struct_scratch.mlir b/test/lit_tests/raising/raise_struct_scratch.mlir new file mode 100644 index 0000000000..ccabcca092 --- /dev/null +++ b/test/lit_tests/raising/raise_struct_scratch.mlir @@ -0,0 +1,67 @@ +// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo --split-input-file | FileCheck %s + +// Struct-element scratch splits into one primitive scratch per field: the +// i32 views address the fields at even/odd offsets, and the whole-pair i64 +// load reads both fields of pair zero at once. +func.func @intpair(%out: memref<256xi32, 1>, %in: memref<256xi32, 1>) { + %scr = memref.alloca() {alignment = 4 : i64} : memref<256x!llvm.struct<"struct.mfem::DevicePair", (i32, i32)>> + %ptr = "enzymexla.memref2pointer"(%scr) : (memref<256x!llvm.struct<"struct.mfem::DevicePair", (i32, i32)>>) -> !llvm.ptr<3> + %c2 = arith.constant 2 : index + affine.parallel (%t) = (0) to (256) { + %view = "enzymexla.pointer2memref"(%ptr) : (!llvm.ptr<3>) -> memref + %v = affine.load %in[%t] : memref<256xi32, 1> + affine.store %v, %view[%t * 2] : memref + %ti = arith.index_castui %t : index to i32 + affine.store %ti, %view[%t * 2 + 1] : memref + "enzymexla.barrier"(%t) : (index) -> () + %wide = "enzymexla.pointer2memref"(%ptr) : (!llvm.ptr<3>) -> memref + %pair0 = affine.load %wide[0] : memref + %lo = arith.trunci %pair0 : i64 to i32 + %t2 = arith.muli %t, %c2 : index + %e = memref.load %view[%t2] : memref + %r = arith.addi %lo, %e : i32 + affine.store %r, %out[%t] : memref<256xi32, 1> + } + return +} + +// CHECK-LABEL: func.func private @intpair_raised( +// CHECK-NOT: llvm.struct +// CHECK: stablehlo.slice + +// ----- + +// Mixed-type pairs keep each field in its own scratch, and the whole-pair +// memcpy between struct-strided geps becomes per-field moves. +func.func @mixedpair(%out: memref<256xf64, 1>, %in: memref<256xf64, 1>) { + %scr = memref.alloca() {alignment = 8 : i64} : memref<256x!llvm.struct<"struct.mfem::DevicePair.0", (f64, i32)>> + %ptr = "enzymexla.memref2pointer"(%scr) : (memref<256x!llvm.struct<"struct.mfem::DevicePair.0", (f64, i32)>>) -> !llvm.ptr<3> + %gp = llvm.addrspacecast %ptr : !llvm.ptr<3> to !llvm.ptr + %c16_i64 = arith.constant 16 : i64 + %c1_i64 = arith.constant 1 : i64 + %c255_i64 = arith.constant 255 : i64 + affine.parallel (%t) = (0) to (256) { + %fview = "enzymexla.pointer2memref"(%ptr) : (!llvm.ptr<3>) -> memref + %iview = "enzymexla.pointer2memref"(%ptr) : (!llvm.ptr<3>) -> memref + %v = affine.load %in[%t] : memref<256xf64, 1> + affine.store %v, %fview[%t * 2] : memref + %ti = arith.index_castui %t : index to i32 + affine.store %ti, %iview[%t * 4 + 2] : memref + "enzymexla.barrier"(%t) : (index) -> () + %tl = arith.index_castui %t : index to i64 + %tn = arith.addi %tl, %c1_i64 : i64 + %ts = arith.andi %tn, %c255_i64 : i64 + %dst = llvm.getelementptr inbounds %gp[%tl] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<16 x i8> + %src = llvm.getelementptr inbounds %gp[%ts] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<16 x i8> + "llvm.intr.memcpy"(%dst, %src, %c16_i64) <{isVolatile = false}> : (!llvm.ptr, !llvm.ptr, i64) -> () + "enzymexla.barrier"(%t) : (index) -> () + %r = affine.load %fview[%t * 2] : memref + affine.store %r, %out[%t] : memref<256xf64, 1> + } + return +} + +// CHECK-LABEL: func.func private @mixedpair_raised( +// CHECK-NOT: llvm.struct +// CHECK: stablehlo.gather +// CHECK: stablehlo.scatter