From cad826a74038617592edd37e282ec2e0e24b1b12 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Fri, 28 Aug 2026 09:12:00 -0500 Subject: [PATCH] While raising: broadcast under-attributed yields to the carried layout A loop-carried value can be yielded with fewer attributed axes than the carried argument (a uniform chain through an index-table gather loses its lane attribution), which made the permutation matching reject the loop. Broadcast the yield up to the carried layout first, and let alignMemoryAccess report incompatibility through an ok flag instead of asserting on accesses it cannot size. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD --- .../jax/Passes/AffineToStableHLORaising.cpp | 59 ++++++++++++++++--- .../affine_to_stablehlo_tridiagonal.mlir | 8 ++- 2 files changed, 57 insertions(+), 10 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp index e3346bd36e..2e1af0612a 100644 --- a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp +++ b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp @@ -494,18 +494,36 @@ affineMapShape(affine::AffineValueMap accessValueMap, ParallelContext pc) { static affine::AffineValueMap alignMemoryAccess(Value &a, affine::AffineValueMap src, Value *bs, ArrayRef dsts, OpBuilder &builder, - ParallelContext pc) { + ParallelContext pc, bool *ok = nullptr) { + // NOTE a default-constructed AffineValueMap holds a null context, so its + // getAffineMap() cannot even be probed; validity is reported through `ok` + // and inputs must be maps the caller actually recorded. + auto fail = [&]() { + if (ok) + *ok = false; + return affine::AffineValueMap(); + }; + if (ok) + *ok = true; + if (!a) + return fail(); + for (unsigned qi = 0; qi < dsts.size(); ++qi) + if (!bs[qi]) + return fail(); // -> tensor<10x1xf32> loaded from (i) -> (i, 0) // -> to tensor<1x10xf32> written as (i) -> (0, i) + // affineMapShape bails to an empty vector on accesses it cannot size (an + // IV with no static range); that must reject the alignment, not assert. SmallVector shapeA = affineMapShape(src, pc); - assert(shapeA.size() == - cast(a.getType()).getShape().size()); + if (shapeA.size() != cast(a.getType()).getShape().size()) + return fail(); SmallVector> shapeBs; for (int i = 0; i < dsts.size(); i++) { shapeBs.push_back(affineMapShape(dsts[i], pc)); - assert(shapeBs[i].size() == - cast(bs[i].getType()).getShape().size()); + if (shapeBs[i].size() != + cast(bs[i].getType()).getShape().size()) + return fail(); } SmallVector outputShape; @@ -617,10 +635,10 @@ alignMemoryAccess(Value &a, affine::AffineValueMap src, Value *bs, static affine::AffineValueMap alignMemoryAccess(Value &a, affine::AffineValueMap src, Value &b, affine::AffineValueMap dst, OpBuilder &builder, - ParallelContext pc) { + ParallelContext pc, bool *ok = nullptr) { Value bs[] = {b}; affine::AffineValueMap dsts[] = {dst}; - auto res = alignMemoryAccess(a, src, bs, dsts, builder, pc); + auto res = alignMemoryAccess(a, src, bs, dsts, builder, pc, ok); b = bs[0]; return res; } @@ -822,6 +840,28 @@ static Block *getRaisedEntryBlock(Operation *op) { return &op->getParentRegion()->front(); } +// A loop-carried value can be yielded with fewer attributed axes than the +// carried argument (a uniform chain through an index-table gather loses its +// lane attribution): broadcast the yield up to the carried layout before +// matching the permutation. +static bool +broadcastYieldToCarried(Value &yielded, Value carried, OpBuilder &builder, + llvm::DenseMap &maps, + ParallelContext &pc) { + if (yielded.getType() == carried.getType()) + return true; + bool ok = true; + Value a = carried; + Value b = yielded; + auto outMap = alignMemoryAccess(a, maps.lookup(carried), b, + maps.lookup(yielded), builder, pc, &ok); + if (!ok || a.getType() != carried.getType()) + return false; + maps[b] = outMap; + yielded = b; + return true; +} + static LogicalResult tryRaisingForOpToStableHLOWhile( affine::AffineForOp forOp, IRMapping &parentMapping, OpBuilder &builder, llvm::DenseMap &maps, ParallelContext pc, @@ -1489,6 +1529,11 @@ static LogicalResult tryRaisingForOpToStableHLOWhile( Value raisedYieldedIterArg = mapping.lookup(yieldedIterArgs); Value raisedIterArg = mapping.lookup(iterArg); + if (!maps.count(raisedYieldedIterArg) || !maps.count(raisedIterArg)) + return failure(); + if (!broadcastYieldToCarried(raisedYieldedIterArg, raisedIterArg, builder, + maps, pc)) + return failure(); auto perm = memoryEquivalentPermutation(maps.lookup(raisedYieldedIterArg), maps.lookup(raisedIterArg)); diff --git a/test/lit_tests/raising/affine_to_stablehlo_tridiagonal.mlir b/test/lit_tests/raising/affine_to_stablehlo_tridiagonal.mlir index a95a8e059b..509b75ce57 100644 --- a/test/lit_tests/raising/affine_to_stablehlo_tridiagonal.mlir +++ b/test/lit_tests/raising/affine_to_stablehlo_tridiagonal.mlir @@ -675,7 +675,9 @@ module { } } -// CHECK: %[[ALIGN1:.+]] = stablehlo.transpose %3191, dims = [0, 2, 3, 1] : (tensor<2x16x16x3xf32>) -> tensor<2x16x3x16xf32> -// CHECK-NEXT: %[[ALIGN2:.+]] = stablehlo.transpose %3508, dims = [0, 2, 3, 1] : (tensor<2x16x16x3xf32>) -> tensor<2x16x3x16xf32> -// CHECK-NEXT: stablehlo.return %{{.+}}, %[[ALIGN1]], %[[ALIGN2]], %{{.+}}, %{{.+}}, %iterArg_491, %iterArg_492, %iterArg_493, %iterArg_494, %iterArg_495, %iterArg_496, %iterArg_497, %iterArg_498, %iterArg_499 : tensor, tensor<2x16x3x16xf32>, tensor<2x16x3x16xf32>, tensor<61x28x46xf32>, tensor<50x18x36xf32>, tensor<61xf32>, tensor<60xf32>, tensor<1x28x46xf32>, tensor<60x28x46xf32>, tensor<60x28x46xf32>, tensor<60x28x46xf32>, tensor, tensor, tensor +// The yielded accumulators re-align (as permuting broadcasts) onto the +// carried layout before the return. +// CHECK: %[[ALIGN1:.+]] = stablehlo.broadcast_in_dim %{{.+}}, dims = [0, 3, 1, 2] : (tensor<2x16x16x3xf32>) -> tensor<2x16x3x16xf32> +// CHECK: %[[ALIGN2:.+]] = stablehlo.broadcast_in_dim %{{.+}}, dims = [0, 3, 1, 2] : (tensor<2x16x16x3xf32>) -> tensor<2x16x3x16xf32> +// CHECK: stablehlo.return %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %iterArg{{.+}} : tensor, tensor<2x16x3x16xf32>, tensor<2x16x3x16xf32>, tensor<61x28x46xf32>, tensor<50x18x36xf32>, tensor<61xf32>, tensor<60xf32>, tensor<1x28x46xf32>, tensor<60x28x46xf32>, tensor<60x28x46xf32>, tensor<60x28x46xf32>, tensor, tensor, tensor // CHECK-NEXT: }