diff --git a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp index e3346bd36e..e1d1845ef5 100644 --- a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp +++ b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp @@ -3660,6 +3660,125 @@ struct AffineToStableHLORaisingPass // of peeling to a serial loop: the axis becomes constant-extent and the // body sits behind an `iv < extent` guard, which the masking machinery // already understands. Barriers over the axis then stay batched no-ops. + // A thread-private array lives inside the lane-batched parallel: every + // lane owns a copy. The raising models buffers as whole tensors, and a + // store whose index map does not involve the lane axes looks uniform, so + // reads would collapse to one lane's value. Give the buffer one leading + // dimension per lane axis and index every access with the lane IVs. + static void privatizeLaneScratch(Operation *root) { + SmallVector allocas; + root->walk([&](memref::AllocaOp a) { allocas.push_back(a); }); + for (auto a : allocas) { + auto par = a->getParentOfType(); + if (!par) + continue; + // Only the nested (thread) parallel batches into lanes; scratch + // directly under the grid parallel is genuinely shared. + if (!par->getParentOfType()) + continue; + if (par.hasMinMaxBounds()) + continue; + auto ranges = par.getConstantRanges(); + if (!ranges) + continue; + int64_t total = 1; + bool ok = true; + for (auto [i, ext] : llvm::enumerate(*ranges)) { + auto lb = getConstant(par.getLowerBoundMap(i)); + if (!lb || *lb != 0 || par.getSteps()[i] != 1 || ext <= 0) { + ok = false; + break; + } + total *= ext; + } + auto MT = cast(a.getType()); + if (!ok || !MT.hasStaticShape() || + total * MT.getNumElements() > (1 << 22)) + continue; + SmallVector accesses; + bool legal = true; + for (Operation *u : a->getUsers()) { + bool isAccess = (isa(u) && + u->getOperand(0) == a.getResult()) || + (isa(u) && + u->getOperand(1) == a.getResult()); + if (!isAccess || !par->isProperAncestor(u)) { + legal = false; + break; + } + accesses.push_back(u); + } + if (!legal) + continue; + + SmallVector newShape(ranges->begin(), ranges->end()); + newShape.append(MT.getShape().begin(), MT.getShape().end()); + OpBuilder b(par); + auto newAlloca = memref::AllocaOp::create( + b, a.getLoc(), + MemRefType::get(newShape, MT.getElementType(), + MemRefLayoutAttrInterface{}, MT.getMemorySpace())); + auto ivs = par.getIVs(); + unsigned K = ivs.size(); + for (Operation *u : accesses) { + if (auto ld = dyn_cast(u)) { + OpBuilder ub(u); + SmallVector idx(ivs.begin(), ivs.end()); + idx.append(ld.getIndices().begin(), ld.getIndices().end()); + Value nl = memref::LoadOp::create(ub, u->getLoc(), newAlloca, idx); + u->getResult(0).replaceAllUsesWith(nl); + u->erase(); + continue; + } + if (auto st = dyn_cast(u)) { + OpBuilder ub(u); + SmallVector idx(ivs.begin(), ivs.end()); + idx.append(st.getIndices().begin(), st.getIndices().end()); + memref::StoreOp::create(ub, u->getLoc(), st.getValue(), newAlloca, + idx); + u->erase(); + continue; + } + AffineMap map; + SmallVector mapOperands; + if (auto ld = dyn_cast(u)) { + map = ld.getAffineMap(); + mapOperands.assign(ld.getMapOperands().begin(), + ld.getMapOperands().end()); + } else { + auto st = cast(u); + map = st.getAffineMap(); + mapOperands.assign(st.getMapOperands().begin(), + st.getMapOperands().end()); + } + unsigned nd = map.getNumDims(); + SmallVector exprs; + for (unsigned k = 0; k < K; ++k) + exprs.push_back(getAffineDimExpr(nd + k, par.getContext())); + for (AffineExpr e : map.getResults()) + exprs.push_back(e); + auto newMap = AffineMap::get(nd + K, map.getNumSymbols(), exprs, + par.getContext()); + SmallVector newOperands(mapOperands.begin(), + mapOperands.begin() + nd); + newOperands.append(ivs.begin(), ivs.end()); + newOperands.append(mapOperands.begin() + nd, mapOperands.end()); + OpBuilder ub(u); + if (auto ld = dyn_cast(u)) { + Value nl = affine::AffineLoadOp::create(ub, u->getLoc(), newAlloca, + newMap, newOperands); + u->getResult(0).replaceAllUsesWith(nl); + } else { + auto st = cast(u); + affine::AffineStoreOp::create(ub, u->getLoc(), st.getValue(), + newAlloca, newMap, newOperands); + } + u->erase(); + } + a.erase(); + } + } + static void boundParallelAxes(Operation *root) { SmallVector worklist; root->walk([&](affine::AffineParallelOp par) { worklist.push_back(par); }); @@ -3901,6 +4020,7 @@ struct AffineToStableHLORaisingPass for (auto func : funcs) { stripAccessMemorySpaceCasts(func); boundParallelAxes(func); + privatizeLaneScratch(func); peelDynamicParallelDims(func); } @@ -3924,6 +4044,7 @@ struct AffineToStableHLORaisingPass for (auto g : gwrap) { stripAccessMemorySpaceCasts(g); boundParallelAxes(g); + privatizeLaneScratch(g); peelDynamicParallelDims(g); } size_t raised_count = 0; diff --git a/test/lit_tests/raising/raise_lane_private_scratch.mlir b/test/lit_tests/raising/raise_lane_private_scratch.mlir new file mode 100644 index 0000000000..eb9bbe9fdc --- /dev/null +++ b/test/lit_tests/raising/raise_lane_private_scratch.mlir @@ -0,0 +1,29 @@ +// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo="prefer_while_raising=false err_if_not_fully_raised=true" | FileCheck %s + +// Scratch allocated inside the lane-batched parallel is private to each +// lane: every lane writes its own values at the same indices, so the +// buffer gains one leading dimension per lane axis. Without the +// privatization all lanes would read lane 0's values. + +// CHECK-LABEL: @lane_raised +// CHECK-NOT: error + +module { + func.func private @lane(%out: memref<64xf64, 1>, %in: memref<64xf64, 1>) { + %two = arith.constant 2.0 : f64 + affine.parallel (%e) = (0) to (4) { + affine.parallel (%t) = (0) to (16) { + %loc = memref.alloca() : memref<2xf64> + %v = affine.load %in[%t] : memref<64xf64, 1> + affine.store %v, %loc[0] : memref<2xf64> + %v2 = arith.mulf %v, %two : f64 + affine.store %v2, %loc[1] : memref<2xf64> + %a = affine.load %loc[0] : memref<2xf64> + %b = affine.load %loc[1] : memref<2xf64> + %s = arith.addf %a, %b : f64 + affine.store %s, %out[%e * 16 + %t] : memref<64xf64, 1> + } + } + return + } +}