Skip to content

Commit 1ffa4ff

Browse files
j2kuncopybara-github
j2kun
authored and
copybara-github
committed
Project out reduced dimension when propagating layout for linalg.reduce
Context: when lowering softmax in a later change rebased on this change, I noticed that the combination of a reduction followed by a broadcast inserted an additional layout conversion that should have been free. In the current lowering for linalg.reduce, the layout propagation does not exactly match the implementation of the kernel. The kernel does a rotate-and-reduce circuit, which has a byproduct of putting the reduced value in all slots of the ciphertext. If the input starts from a layout like R(i₀,ct,slot) ; ct = 0 and (slot - i₀) mod 8 = 0 and 0 ≤ i₀ ≤ 7 and 0 ≤ slot ≤ 2047 then the output layout should be R(ct,slot) ; ct = 0 and 0 ≤ slot ≤ 2047 This represents a layout where the input scalar is present in every slot. However, the way the current propagation works, it sets the i₀ to zero first, which produces R(ct,slot) ; ct = 0 and slot mod 8 = 0 and 0 ≤ slot ≤ 2047 This describes a layout where the scalar is only in slots whose index is a multiple of 8. This change switches from zero-setting to projection, and updates the test to check more explicitly for the desired output layout. PiperOrigin-RevId: 952818293
1 parent f5d3838 commit 1ffa4ff

14 files changed

Lines changed: 563 additions & 3 deletions

File tree

‎lib/Transforms/LayoutPropagation/Utils.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -107,11 +107,11 @@ LayoutAttr convertLayoutForReduce(LayoutAttr inputLayout,
107107

108108
auto offset = clonedRelation->getVarKindOffset(presburger::VarKind::Domain);
109109
for (int dim : llvm::reverse(dimsToReduce)) {
110-
// Set the dim to reduce equal to 0.
110+
// Project out the reduced dimension.
111111
auto dimIndex = offset + dim;
112112
assert(clonedRelation->getVarKindAt(dimIndex) ==
113113
presburger::VarKind::Domain);
114-
clonedRelation->setAndEliminate(dimIndex, 0);
114+
clonedRelation->projectOut(dimIndex, 1);
115115
}
116116

117117
MLIRContext* context = inputLayout.getContext();

‎lib/Transforms/SoftmaxToCgfSoftmax/BUILD‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
load("@heir//lib/Transforms:transforms.bzl", "add_heir_transforms")
2+
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
3+
load("@rules_cc//cc:cc_library.bzl", "cc_library")
4+
5+
package(
6+
default_applicable_licenses = ["@heir//:license"],
7+
default_visibility = ["//visibility:public"],
8+
)
9+
10+
cc_library(
11+
name = "SoftmaxToCgfSoftmax",
12+
srcs = ["SoftmaxToCgfSoftmax.cpp"],
13+
hdrs = ["SoftmaxToCgfSoftmax.h"],
14+
deps = [
15+
":pass_inc_gen",
16+
"@heir//lib/Dialect/MathExt/IR:Dialect",
17+
"@llvm-project//mlir:ArithDialect",
18+
"@llvm-project//mlir:DialectUtils",
19+
"@llvm-project//mlir:IR",
20+
"@llvm-project//mlir:LinalgDialect",
21+
"@llvm-project//mlir:MathDialect",
22+
"@llvm-project//mlir:Pass",
23+
"@llvm-project//mlir:Support",
24+
"@llvm-project//mlir:TensorDialect",
25+
"@llvm-project//mlir:TransformUtils",
26+
],
27+
)
28+
29+
add_heir_transforms(
30+
generated_target_name = "pass_inc_gen",
31+
pass_name = "SoftmaxToCgfSoftmax",
32+
td_file = "SoftmaxToCgfSoftmax.td",
33+
)
34+
35+
cc_binary(
36+
name = "cgf_softmax_fuzz",
37+
srcs = ["cgf_softmax_fuzz.cpp"],
38+
deps = [
39+
"@com_google_absl//absl/random",
40+
],
41+
)

‎lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.cpp‎

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
#include "lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.h"
2+
3+
#include <algorithm>
4+
#include <cmath>
5+
#include <cstdint>
6+
#include <utility>
7+
8+
#include "lib/Dialect/MathExt/IR/MathExtOps.h"
9+
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
10+
#include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h" // from @llvm-project
11+
#include "mlir/include/mlir/Dialect/Math/IR/Math.h" // from @llvm-project
12+
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
13+
#include "mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h" // from @llvm-project
14+
#include "mlir/include/mlir/IR/AffineMap.h" // from @llvm-project
15+
#include "mlir/include/mlir/IR/BuiltinAttributes.h" // from @llvm-project
16+
#include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project
17+
#include "mlir/include/mlir/IR/Location.h" // from @llvm-project
18+
#include "mlir/include/mlir/IR/PatternMatch.h" // from @llvm-project
19+
#include "mlir/include/mlir/IR/TypeRange.h" // from @llvm-project
20+
#include "mlir/include/mlir/IR/Types.h" // from @llvm-project
21+
#include "mlir/include/mlir/IR/Value.h" // from @llvm-project
22+
#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project
23+
#include "mlir/include/mlir/Transforms/WalkPatternRewriteDriver.h" // from @llvm-project
24+
25+
namespace mlir {
26+
namespace heir {
27+
28+
#define GEN_PASS_DEF_SOFTMAXTOCGFSOFTMAX
29+
#include "lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.h.inc"
30+
31+
namespace {
32+
33+
// Helper to create a linalg.reduce sum operation.
34+
// Returns the reduced tensor.
35+
Value createSumReduction(PatternRewriter& rewriter, Location loc, Value input,
36+
Type elemType, int64_t reductionDim) {
37+
auto inputType = cast<RankedTensorType>(input.getType());
38+
auto inputShape = inputType.getShape();
39+
SmallVector<int64_t> outputShape;
40+
for (int i = 0; i < inputType.getRank(); ++i) {
41+
if (i != reductionDim) {
42+
outputShape.push_back(inputShape[i]);
43+
}
44+
}
45+
auto outputType = RankedTensorType::get(outputShape, elemType);
46+
auto splatAttr =
47+
DenseElementsAttr::get(outputType, rewriter.getFloatAttr(elemType, 0.0));
48+
Value filled = arith::ConstantOp::create(rewriter, loc, splatAttr);
49+
50+
SmallVector<int64_t> dimensions = {reductionDim};
51+
auto reduceOp =
52+
linalg::ReduceOp::create(rewriter, loc,
53+
/*resultTypes=*/TypeRange{filled.getType()},
54+
/*inputs=*/ValueRange{input},
55+
/*inits=*/ValueRange{filled},
56+
/*dimensions=*/dimensions);
57+
58+
{
59+
OpBuilder::InsertionGuard guard(rewriter);
60+
Block* body =
61+
rewriter.createBlock(&reduceOp.getRegion(), reduceOp.getRegion().end(),
62+
TypeRange{elemType, elemType}, {loc, loc});
63+
Value add = arith::AddFOp::create(rewriter, loc, body->getArgument(0),
64+
body->getArgument(1));
65+
linalg::YieldOp::create(rewriter, loc, add);
66+
}
67+
return reduceOp.getResult(0);
68+
}
69+
70+
struct SoftmaxToCgfSoftmaxPattern
71+
: public OpRewritePattern<math_ext::SoftmaxOp> {
72+
using OpRewritePattern<math_ext::SoftmaxOp>::OpRewritePattern;
73+
74+
LogicalResult matchAndRewrite(math_ext::SoftmaxOp op,
75+
PatternRewriter& rewriter) const override {
76+
Location loc = op.getLoc();
77+
Value input = op.getValue();
78+
auto inputType = cast<RankedTensorType>(input.getType());
79+
assert(inputType.hasStaticShape() && "only static shapes are supported");
80+
int64_t rank = inputType.getRank();
81+
assert((rank == 1 || rank == 2) && "only 1D and 2D tensors are supported");
82+
83+
Type elemType = inputType.getElementType();
84+
auto inputShape = inputType.getShape();
85+
int64_t n = inputShape[rank - 1];
86+
double n_double = static_cast<double>(n);
87+
88+
Value invNConst = arith::ConstantOp::create(
89+
rewriter, loc, rewriter.getFloatAttr(elemType, 1.0 / n_double));
90+
Value halfConst = arith::ConstantOp::create(
91+
rewriter, loc, rewriter.getFloatAttr(elemType, 0.5));
92+
Value lnNConst = arith::ConstantOp::create(
93+
rewriter, loc, rewriter.getFloatAttr(elemType, std::log(n_double)));
94+
95+
int64_t reductionDim = rank - 1;
96+
SmallVector<int64_t> reductionShape(inputShape.begin(), inputShape.end());
97+
reductionShape.erase(reductionShape.begin() + reductionDim);
98+
auto reductionType = RankedTensorType::get(reductionShape, elemType);
99+
100+
// 1. Compute mean (mu)
101+
Value sum =
102+
createSumReduction(rewriter, loc, input, elemType, reductionDim);
103+
Value invNConstSplat =
104+
tensor::SplatOp::create(rewriter, loc, reductionType, invNConst);
105+
Value mu = arith::MulFOp::create(rewriter, loc, sum, invNConstSplat);
106+
107+
// 2. Compute variance (sigma^2)
108+
Value initTensor =
109+
tensor::EmptyOp::create(rewriter, loc, inputShape, elemType);
110+
111+
// Broadcast mu along the reduced dimension
112+
Value muBroadcast =
113+
linalg::BroadcastOp::create(rewriter, loc, mu, initTensor,
114+
ArrayRef<int64_t>{reductionDim})
115+
.getResults()[0];
116+
Value diff = arith::SubFOp::create(rewriter, loc, input, muBroadcast);
117+
Value diffSq = arith::MulFOp::create(rewriter, loc, diff, diff);
118+
119+
Value sumDiffSq =
120+
createSumReduction(rewriter, loc, diffSq, elemType, reductionDim);
121+
Value sigmaSq =
122+
arith::MulFOp::create(rewriter, loc, sumDiffSq, invNConstSplat);
123+
124+
// 3. Compute shift S = mu + sigma_sq / 2 + ln(n)
125+
Value halfSplat =
126+
tensor::SplatOp::create(rewriter, loc, reductionType, halfConst);
127+
Value lnNSplat =
128+
tensor::SplatOp::create(rewriter, loc, reductionType, lnNConst);
129+
Value halfSigmaSq =
130+
arith::MulFOp::create(rewriter, loc, sigmaSq, halfSplat);
131+
Value muPlusHalfSigmaSq =
132+
arith::AddFOp::create(rewriter, loc, mu, halfSigmaSq);
133+
Value shift =
134+
arith::AddFOp::create(rewriter, loc, muPlusHalfSigmaSq, lnNSplat);
135+
136+
// 4. Shift inputs and apply exp: result = exp(input - shift)
137+
double L_val =
138+
op->hasAttr("domain_lower")
139+
? cast<FloatAttr>(op->getAttr("domain_lower")).getValueAsDouble()
140+
: -1.0;
141+
double U_val =
142+
op->hasAttr("domain_upper")
143+
? cast<FloatAttr>(op->getAttr("domain_upper")).getValueAsDouble()
144+
: 1.0;
145+
double est_lower =
146+
L_val -
147+
(U_val + (U_val - L_val) * (U_val - L_val) / 8.0 + std::log(n_double));
148+
double safe_lower = std::max(est_lower, -16.0);
149+
150+
Value shiftBroadcast =
151+
linalg::BroadcastOp::create(rewriter, loc, shift, initTensor,
152+
ArrayRef<int64_t>{reductionDim})
153+
.getResults()[0];
154+
Value shiftedInput =
155+
arith::SubFOp::create(rewriter, loc, input, shiftBroadcast);
156+
auto expOp = math::ExpOp::create(rewriter, loc, shiftedInput);
157+
expOp->setAttr("domain_lower", rewriter.getF64FloatAttr(safe_lower));
158+
expOp->setAttr("domain_upper", rewriter.getF64FloatAttr(0.5));
159+
160+
rewriter.replaceOp(op, expOp->getResults());
161+
return success();
162+
}
163+
};
164+
165+
} // namespace
166+
167+
struct SoftmaxToCgfSoftmaxPass
168+
: public impl::SoftmaxToCgfSoftmaxBase<SoftmaxToCgfSoftmaxPass> {
169+
void runOnOperation() override {
170+
MLIRContext* context = &getContext();
171+
Operation* op = getOperation();
172+
173+
// Pre-check for errors/warnings on domain width.
174+
WalkResult walkResult = op->walk([](math_ext::SoftmaxOp softmaxOp) {
175+
auto lowerAttr =
176+
dyn_cast_or_null<FloatAttr>(softmaxOp->getAttr("domain_lower"));
177+
auto upperAttr =
178+
dyn_cast_or_null<FloatAttr>(softmaxOp->getAttr("domain_upper"));
179+
if (lowerAttr && upperAttr) {
180+
double L = lowerAttr.getValueAsDouble();
181+
double U = upperAttr.getValueAsDouble();
182+
double width = U - L;
183+
if (width > 4.0) {
184+
softmaxOp->emitOpError()
185+
<< "input domain width (" << width
186+
<< ") exceeds the maximum safe limit (4.0) for CGF-softmax "
187+
"approximation";
188+
return WalkResult::interrupt();
189+
} else if (width > 2.0) {
190+
softmaxOp->emitWarning()
191+
<< "input domain width (" << width
192+
<< ") exceeds the recommended safe limit (2.0) for CGF-softmax "
193+
"approximation. Accuracy may degrade.";
194+
}
195+
}
196+
return WalkResult::advance();
197+
});
198+
199+
if (walkResult.wasInterrupted()) {
200+
signalPassFailure();
201+
return;
202+
}
203+
204+
RewritePatternSet patterns(context);
205+
patterns.add<SoftmaxToCgfSoftmaxPattern>(context);
206+
207+
walkAndApplyPatterns(op, std::move(patterns));
208+
}
209+
};
210+
211+
} // namespace heir
212+
} // namespace mlir

‎lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
#ifndef LIB_TRANSFORMS_SOFTMAXTOCGFSOFTMAX_SOFTMAXTOCGFSOFTMAX_H_
2+
#define LIB_TRANSFORMS_SOFTMAXTOCGFSOFTMAX_SOFTMAXTOCGFSOFTMAX_H_
3+
4+
#include "mlir/include/mlir/Pass/Pass.h" // from @llvm-project
5+
6+
namespace mlir {
7+
namespace heir {
8+
9+
#define GEN_PASS_DECL
10+
#include "lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.h.inc"
11+
12+
#define GEN_PASS_REGISTRATION
13+
#include "lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.h.inc"
14+
15+
} // namespace heir
16+
} // namespace mlir
17+
18+
#endif // LIB_TRANSFORMS_SOFTMAXTOCGFSOFTMAX_SOFTMAXTOCGFSOFTMAX_H_

‎lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.td‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
#ifndef LIB_TRANSFORMS_SOFTMAXTOCGFSOFTMAX_SOFTMAXTOCGFSOFTMAX_TD_
2+
#define LIB_TRANSFORMS_SOFTMAXTOCGFSOFTMAX_SOFTMAXTOCGFSOFTMAX_TD_
3+
4+
include "mlir/Pass/PassBase.td"
5+
6+
def SoftmaxToCgfSoftmax : Pass<"softmax-to-cgf-softmax"> {
7+
let summary = "Lower math_ext.softmax to CGF-softmax approximation";
8+
let description = [{
9+
Lowers `math_ext.softmax` to a sequence of operations implementing
10+
the second-order CGF-softmax approximation. This eliminates
11+
division and max operations.
12+
13+
For details on the approximation, see: https://arxiv.org/abs/2602.01621
14+
}];
15+
let dependentDialects = [
16+
"mlir::arith::ArithDialect",
17+
"mlir::math::MathDialect",
18+
"mlir::linalg::LinalgDialect",
19+
"mlir::tensor::TensorDialect",
20+
"mlir::heir::math_ext::MathExtDialect"
21+
];
22+
}
23+
24+
#endif // LIB_TRANSFORMS_SOFTMAXTOCGFSOFTMAX_SOFTMAXTOCGFSOFTMAX_TD_

0 commit comments

Comments
 (0)