Skip to content

Commit c1ebefd

Browse files
d0kd0k
committed
[mlir] Make polynomial approximation emit std instead of LLVM ops
This is a bit cleaner and removes issues with 2d vectors. It also has a big impact on constant folding, hence the test changes. Differential Revision: https://reviews.llvm.org/D107896
1 parent a0d8a08 commit c1ebefd

File tree

3 files changed

+62
-77
lines changed

3 files changed

+62
-77
lines changed

‎mlir/lib/Dialect/Math/Transforms/CMakeLists.txt

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ add_mlir_dialect_library(MLIRMathTransforms
88

99
LINK_LIBS PUBLIC
1010
MLIRIR
11-
MLIRLLVMIR
1211
MLIRMath
1312
MLIRPass
1413
MLIRStandard

‎mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@
1111
//
1212
//===----------------------------------------------------------------------===//
1313

14-
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
15-
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
1614
#include "mlir/Dialect/Math/IR/Math.h"
1715
#include "mlir/Dialect/Math/Transforms/Passes.h"
1816
#include "mlir/Dialect/Vector/VectorOps.h"
@@ -96,7 +94,7 @@ static Value i32Cst(ImplicitLocOpBuilder &builder, int32_t value) {
9694

9795
static Value f32FromBits(ImplicitLocOpBuilder &builder, uint32_t bits) {
9896
Value i32Value = i32Cst(builder, static_cast<int32_t>(bits));
99-
return builder.create<LLVM::BitcastOp>(builder.getF32Type(), i32Value);
97+
return builder.create<BitcastOp>(builder.getF32Type(), i32Value);
10098
}
10199

102100
//----------------------------------------------------------------------------//
@@ -139,20 +137,19 @@ static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg,
139137
Value cstInvMantMask = f32FromBits(builder, ~0x7f800000u);
140138

141139
// Bitcast to i32 for bitwise operations.
142-
Value i32Half = builder.create<LLVM::BitcastOp>(i32, cstHalf);
143-
Value i32InvMantMask = builder.create<LLVM::BitcastOp>(i32, cstInvMantMask);
144-
Value i32Arg = builder.create<LLVM::BitcastOp>(i32Vec, arg);
140+
Value i32Half = builder.create<BitcastOp>(i32, cstHalf);
141+
Value i32InvMantMask = builder.create<BitcastOp>(i32, cstInvMantMask);
142+
Value i32Arg = builder.create<BitcastOp>(i32Vec, arg);
145143

146144
// Compute normalized fraction.
147-
Value tmp0 = builder.create<LLVM::AndOp>(i32Arg, bcast(i32InvMantMask));
148-
Value tmp1 = builder.create<LLVM::OrOp>(tmp0, bcast(i32Half));
149-
Value normalizedFraction = builder.create<LLVM::BitcastOp>(f32Vec, tmp1);
145+
Value tmp0 = builder.create<AndOp>(i32Arg, bcast(i32InvMantMask));
146+
Value tmp1 = builder.create<OrOp>(tmp0, bcast(i32Half));
147+
Value normalizedFraction = builder.create<BitcastOp>(f32Vec, tmp1);
150148

151149
// Compute exponent.
152150
Value arg0 = is_positive ? arg : builder.create<AbsFOp>(arg);
153151
Value biasedExponentBits = builder.create<UnsignedShiftRightOp>(
154-
builder.create<LLVM::BitcastOp>(i32Vec, arg0),
155-
bcast(i32Cst(builder, 23)));
152+
builder.create<BitcastOp>(i32Vec, arg0), bcast(i32Cst(builder, 23)));
156153
Value biasedExponent = builder.create<SIToFPOp>(f32Vec, biasedExponentBits);
157154
Value exponent = builder.create<SubFOp>(biasedExponent, bcast(cst126f));
158155

@@ -178,7 +175,7 @@ static Value exp2I32(ImplicitLocOpBuilder &builder, Value arg) {
178175
Value biasedArg = builder.create<AddIOp>(arg, bias);
179176
Value exp2ValueInt =
180177
builder.create<ShiftLeftOp>(biasedArg, exponetBitLocation);
181-
Value exp2ValueF32 = builder.create<LLVM::BitcastOp>(f32Vec, exp2ValueInt);
178+
Value exp2ValueF32 = builder.create<BitcastOp>(f32Vec, exp2ValueInt);
182179

183180
return exp2ValueF32;
184181
}
@@ -454,8 +451,8 @@ Log1pApproximation::matchAndRewrite(math::Log1pOp op,
454451
Value uInf = builder.create<CmpFOp>(CmpFPredicate::OEQ, u, logU);
455452
Value logLarge = builder.create<MulFOp>(
456453
x, builder.create<DivFOp>(logU, builder.create<SubFOp>(u, cstOne)));
457-
Value approximation = builder.create<SelectOp>(
458-
builder.create<LLVM::OrOp>(uSmall, uInf), x, logLarge);
454+
Value approximation =
455+
builder.create<SelectOp>(builder.create<OrOp>(uSmall, uInf), x, logLarge);
459456
rewriter.replaceOp(op, approximation);
460457
return success();
461458
}

‎mlir/test/Dialect/Math/polynomial-approximation.mlir

Lines changed: 51 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
// CHECK: %[[VAL_26:.*]] = fptosi %[[VAL_16]] : f32 to i32
3434
// CHECK: %[[VAL_27:.*]] = addi %[[VAL_26]], %[[VAL_13]] : i32
3535
// CHECK: %[[VAL_28:.*]] = shift_left %[[VAL_27]], %[[VAL_8]] : i32
36-
// CHECK: %[[VAL_29:.*]] = llvm.bitcast %[[VAL_28]] : i32 to f32
36+
// CHECK: %[[VAL_29:.*]] = bitcast %[[VAL_28]] : i32 to f32
3737
// CHECK: %[[VAL_30:.*]] = mulf %[[VAL_25]], %[[VAL_29]] : f32
3838
// CHECK: %[[VAL_31:.*]] = cmpi sle, %[[VAL_26]], %[[VAL_13]] : i32
3939
// CHECK: %[[VAL_32:.*]] = cmpi sge, %[[VAL_26]], %[[VAL_14]] : i32
@@ -97,10 +97,9 @@ func @expm1_scalar(%arg0: f32) -> f32 {
9797
// CHECK-NOT: exp
9898
// CHECK-COUNT-3: select
9999
// CHECK-NOT: log
100-
// CHECK-COUNT-6: vector.broadcast
101100
// CHECK-COUNT-5: select
102101
// CHECK-NOT: expm1
103-
// CHECK-COUNT-2: select
102+
// CHECK-COUNT-3: select
104103
// CHECK: %[[VAL_115:.*]] = select
105104
// CHECK: return %[[VAL_115]] : vector<8xf32>
106105
// CHECK: }
@@ -110,14 +109,14 @@ func @expm1_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
110109
}
111110

112111
// CHECK-LABEL: func @log_scalar(
113-
// CHECK-SAME: %[[X:.*]]: f32) -> f32 {
112+
// CHECK-SAME: %[[X:.*]]: f32) -> f32 {
114113
// CHECK: %[[VAL_1:.*]] = constant 0.000000e+00 : f32
115114
// CHECK: %[[VAL_2:.*]] = constant 1.000000e+00 : f32
116115
// CHECK: %[[VAL_3:.*]] = constant -5.000000e-01 : f32
117-
// CHECK: %[[VAL_4:.*]] = constant 8388608 : i32
118-
// CHECK: %[[VAL_5:.*]] = constant -8388608 : i32
119-
// CHECK: %[[VAL_6:.*]] = constant 2139095040 : i32
120-
// CHECK: %[[VAL_7:.*]] = constant 2143289344 : i32
116+
// CHECK: %[[VAL_4:.*]] = constant 1.17549435E-38 : f32
117+
// CHECK: %[[VAL_5:.*]] = constant 0xFF800000 : f32
118+
// CHECK: %[[VAL_6:.*]] = constant 0x7F800000 : f32
119+
// CHECK: %[[VAL_7:.*]] = constant 0x7FC00000 : f32
121120
// CHECK: %[[VAL_8:.*]] = constant 0.707106769 : f32
122121
// CHECK: %[[VAL_9:.*]] = constant 0.0703768358 : f32
123122
// CHECK: %[[VAL_10:.*]] = constant -0.115146101 : f32
@@ -129,55 +128,48 @@ func @expm1_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
129128
// CHECK: %[[VAL_16:.*]] = constant -0.24999994 : f32
130129
// CHECK: %[[VAL_17:.*]] = constant 0.333333313 : f32
131130
// CHECK: %[[VAL_18:.*]] = constant 1.260000e+02 : f32
132-
// CHECK: %[[VAL_19:.*]] = constant 5.000000e-01 : f32
133-
// CHECK: %[[VAL_20:.*]] = constant -2139095041 : i32
131+
// CHECK: %[[VAL_19:.*]] = constant -2139095041 : i32
132+
// CHECK: %[[VAL_20:.*]] = constant 1056964608 : i32
134133
// CHECK: %[[VAL_21:.*]] = constant 23 : i32
135-
// CHECK: %[[CST_LN2:.*]] = constant 0.693147182 : f32
136-
// CHECK: %[[VAL_23:.*]] = llvm.bitcast %[[VAL_4]] : i32 to f32
137-
// CHECK: %[[VAL_24:.*]] = llvm.bitcast %[[VAL_5]] : i32 to f32
138-
// CHECK: %[[VAL_25:.*]] = llvm.bitcast %[[VAL_6]] : i32 to f32
139-
// CHECK: %[[VAL_26:.*]] = llvm.bitcast %[[VAL_7]] : i32 to f32
140-
// CHECK: %[[VAL_27:.*]] = cmpf ogt, %[[X]], %[[VAL_23]] : f32
141-
// CHECK: %[[VAL_28:.*]] = select %[[VAL_27]], %[[X]], %[[VAL_23]] : f32
134+
// CHECK: %[[VAL_22:.*]] = constant 0.693147182 : f32
135+
// CHECK: %[[VAL_23:.*]] = cmpf ogt, %[[X]], %[[VAL_4]] : f32
136+
// CHECK: %[[VAL_24:.*]] = select %[[VAL_23]], %[[X]], %[[VAL_4]] : f32
142137
// CHECK-NOT: frexp
143-
// CHECK: %[[VAL_29:.*]] = llvm.bitcast %[[VAL_20]] : i32 to f32
144-
// CHECK: %[[VAL_30:.*]] = llvm.bitcast %[[VAL_19]] : f32 to i32
145-
// CHECK: %[[VAL_31:.*]] = llvm.bitcast %[[VAL_29]] : f32 to i32
146-
// CHECK: %[[VAL_32:.*]] = llvm.bitcast %[[VAL_28]] : f32 to i32
147-
// CHECK: %[[VAL_33:.*]] = llvm.and %[[VAL_32]], %[[VAL_31]] : i32
148-
// CHECK: %[[VAL_34:.*]] = llvm.or %[[VAL_33]], %[[VAL_30]] : i32
149-
// CHECK: %[[VAL_35:.*]] = llvm.bitcast %[[VAL_34]] : i32 to f32
150-
// CHECK: %[[VAL_36:.*]] = llvm.bitcast %[[VAL_28]] : f32 to i32
151-
// CHECK: %[[VAL_37:.*]] = shift_right_unsigned %[[VAL_36]], %[[VAL_21]] : i32
152-
// CHECK: %[[FREXP_X:.*]] = sitofp %[[VAL_37]] : i32 to f32
153-
// CHECK: %[[VAL_39:.*]] = subf %[[FREXP_X]], %[[VAL_18]] : f32
154-
// CHECK: %[[VAL_40:.*]] = cmpf olt, %[[VAL_35]], %[[VAL_8]] : f32
155-
// CHECK: %[[VAL_41:.*]] = select %[[VAL_40]], %[[VAL_35]], %[[VAL_1]] : f32
156-
// CHECK: %[[VAL_42:.*]] = subf %[[VAL_35]], %[[VAL_2]] : f32
157-
// CHECK: %[[VAL_43:.*]] = select %[[VAL_40]], %[[VAL_2]], %[[VAL_1]] : f32
158-
// CHECK: %[[VAL_44:.*]] = subf %[[VAL_39]], %[[VAL_43]] : f32
159-
// CHECK: %[[VAL_45:.*]] = addf %[[VAL_42]], %[[VAL_41]] : f32
160-
// CHECK: %[[VAL_46:.*]] = mulf %[[VAL_45]], %[[VAL_45]] : f32
161-
// CHECK: %[[VAL_47:.*]] = mulf %[[VAL_46]], %[[VAL_45]] : f32
162-
// CHECK: %[[VAL_48:.*]] = fmaf %[[VAL_9]], %[[VAL_45]], %[[VAL_10]] : f32
163-
// CHECK: %[[VAL_49:.*]] = fmaf %[[VAL_12]], %[[VAL_45]], %[[VAL_13]] : f32
164-
// CHECK: %[[VAL_50:.*]] = fmaf %[[VAL_15]], %[[VAL_45]], %[[VAL_16]] : f32
165-
// CHECK: %[[VAL_51:.*]] = fmaf %[[VAL_48]], %[[VAL_45]], %[[VAL_11]] : f32
166-
// CHECK: %[[VAL_52:.*]] = fmaf %[[VAL_49]], %[[VAL_45]], %[[VAL_14]] : f32
167-
// CHECK: %[[VAL_53:.*]] = fmaf %[[VAL_50]], %[[VAL_45]], %[[VAL_17]] : f32
168-
// CHECK: %[[VAL_54:.*]] = fmaf %[[VAL_51]], %[[VAL_47]], %[[VAL_52]] : f32
169-
// CHECK: %[[VAL_55:.*]] = fmaf %[[VAL_54]], %[[VAL_47]], %[[VAL_53]] : f32
170-
// CHECK: %[[VAL_56:.*]] = mulf %[[VAL_55]], %[[VAL_47]] : f32
171-
// CHECK: %[[VAL_57:.*]] = fmaf %[[VAL_3]], %[[VAL_46]], %[[VAL_56]] : f32
172-
// CHECK: %[[VAL_58:.*]] = addf %[[VAL_45]], %[[VAL_57]] : f32
173-
// CHECK: %[[VAL_59:.*]] = fmaf %[[VAL_44]], %[[CST_LN2]], %[[VAL_58]] : f32
174-
// CHECK: %[[VAL_60:.*]] = cmpf ult, %[[X]], %[[VAL_1]] : f32
175-
// CHECK: %[[VAL_61:.*]] = cmpf oeq, %[[X]], %[[VAL_1]] : f32
176-
// CHECK: %[[VAL_62:.*]] = cmpf oeq, %[[X]], %[[VAL_25]] : f32
177-
// CHECK: %[[VAL_63:.*]] = select %[[VAL_62]], %[[VAL_25]], %[[VAL_59]] : f32
178-
// CHECK: %[[VAL_64:.*]] = select %[[VAL_60]], %[[VAL_26]], %[[VAL_63]] : f32
179-
// CHECK: %[[VAL_65:.*]] = select %[[VAL_61]], %[[VAL_24]], %[[VAL_64]] : f32
180-
// CHECK: return %[[VAL_65]] : f32
138+
// CHECK: %[[VAL_25:.*]] = bitcast %[[VAL_24]] : f32 to i32
139+
// CHECK: %[[VAL_26:.*]] = and %[[VAL_25]], %[[VAL_19]] : i32
140+
// CHECK: %[[VAL_27:.*]] = or %[[VAL_26]], %[[VAL_20]] : i32
141+
// CHECK: %[[VAL_28:.*]] = bitcast %[[VAL_27]] : i32 to f32
142+
// CHECK: %[[VAL_29:.*]] = bitcast %[[VAL_24]] : f32 to i32
143+
// CHECK: %[[VAL_30:.*]] = shift_right_unsigned %[[VAL_29]], %[[VAL_21]] : i32
144+
// CHECK: %[[VAL_31:.*]] = sitofp %[[VAL_30]] : i32 to f32
145+
// CHECK: %[[VAL_32:.*]] = subf %[[VAL_31]], %[[VAL_18]] : f32
146+
// CHECK: %[[VAL_33:.*]] = cmpf olt, %[[VAL_28]], %[[VAL_8]] : f32
147+
// CHECK: %[[VAL_34:.*]] = select %[[VAL_33]], %[[VAL_28]], %[[VAL_1]] : f32
148+
// CHECK: %[[VAL_35:.*]] = subf %[[VAL_28]], %[[VAL_2]] : f32
149+
// CHECK: %[[VAL_36:.*]] = select %[[VAL_33]], %[[VAL_2]], %[[VAL_1]] : f32
150+
// CHECK: %[[VAL_37:.*]] = subf %[[VAL_32]], %[[VAL_36]] : f32
151+
// CHECK: %[[VAL_38:.*]] = addf %[[VAL_35]], %[[VAL_34]] : f32
152+
// CHECK: %[[VAL_39:.*]] = mulf %[[VAL_38]], %[[VAL_38]] : f32
153+
// CHECK: %[[VAL_40:.*]] = mulf %[[VAL_39]], %[[VAL_38]] : f32
154+
// CHECK: %[[VAL_41:.*]] = fmaf %[[VAL_9]], %[[VAL_38]], %[[VAL_10]] : f32
155+
// CHECK: %[[VAL_42:.*]] = fmaf %[[VAL_12]], %[[VAL_38]], %[[VAL_13]] : f32
156+
// CHECK: %[[VAL_43:.*]] = fmaf %[[VAL_15]], %[[VAL_38]], %[[VAL_16]] : f32
157+
// CHECK: %[[VAL_44:.*]] = fmaf %[[VAL_41]], %[[VAL_38]], %[[VAL_11]] : f32
158+
// CHECK: %[[VAL_45:.*]] = fmaf %[[VAL_42]], %[[VAL_38]], %[[VAL_14]] : f32
159+
// CHECK: %[[VAL_46:.*]] = fmaf %[[VAL_43]], %[[VAL_38]], %[[VAL_17]] : f32
160+
// CHECK: %[[VAL_47:.*]] = fmaf %[[VAL_44]], %[[VAL_40]], %[[VAL_45]] : f32
161+
// CHECK: %[[VAL_48:.*]] = fmaf %[[VAL_47]], %[[VAL_40]], %[[VAL_46]] : f32
162+
// CHECK: %[[VAL_49:.*]] = mulf %[[VAL_48]], %[[VAL_40]] : f32
163+
// CHECK: %[[VAL_50:.*]] = fmaf %[[VAL_3]], %[[VAL_39]], %[[VAL_49]] : f32
164+
// CHECK: %[[VAL_51:.*]] = addf %[[VAL_38]], %[[VAL_50]] : f32
165+
// CHECK: %[[VAL_52:.*]] = fmaf %[[VAL_37]], %[[VAL_22]], %[[VAL_51]] : f32
166+
// CHECK: %[[VAL_53:.*]] = cmpf ult, %[[X]], %[[VAL_1]] : f32
167+
// CHECK: %[[VAL_54:.*]] = cmpf oeq, %[[X]], %[[VAL_1]] : f32
168+
// CHECK: %[[VAL_55:.*]] = cmpf oeq, %[[X]], %[[VAL_6]] : f32
169+
// CHECK: %[[VAL_56:.*]] = select %[[VAL_55]], %[[VAL_6]], %[[VAL_52]] : f32
170+
// CHECK: %[[VAL_57:.*]] = select %[[VAL_53]], %[[VAL_7]], %[[VAL_56]] : f32
171+
// CHECK: %[[VAL_58:.*]] = select %[[VAL_54]], %[[VAL_5]], %[[VAL_57]] : f32
172+
// CHECK: return %[[VAL_58]] : f32
181173
// CHECK: }
182174
func @log_scalar(%arg0: f32) -> f32 {
183175
%0 = math.log %arg0 : f32
@@ -187,8 +179,7 @@ func @log_scalar(%arg0: f32) -> f32 {
187179
// CHECK-LABEL: func @log_vector(
188180
// CHECK-SAME: %[[VAL_0:.*]]: vector<8xf32>) -> vector<8xf32> {
189181
// CHECK: %[[CST_LN2:.*]] = constant dense<0.693147182> : vector<8xf32>
190-
// CHECK-COUNT-6: vector.broadcast
191-
// CHECK-COUNT-4: select
182+
// CHECK-COUNT-5: select
192183
// CHECK: %[[VAL_71:.*]] = select
193184
// CHECK: return %[[VAL_71]] : vector<8xf32>
194185
// CHECK: }
@@ -212,8 +203,7 @@ func @log2_scalar(%arg0: f32) -> f32 {
212203
// CHECK-LABEL: func @log2_vector(
213204
// CHECK-SAME: %[[VAL_0:.*]]: vector<8xf32>) -> vector<8xf32> {
214205
// CHECK: %[[CST_LOG2E:.*]] = constant dense<1.44269502> : vector<8xf32>
215-
// CHECK-COUNT-6: vector.broadcast
216-
// CHECK-COUNT-4: select
206+
// CHECK-COUNT-5: select
217207
// CHECK: %[[VAL_71:.*]] = select
218208
// CHECK: return %[[VAL_71]] : vector<8xf32>
219209
// CHECK: }
@@ -234,7 +224,7 @@ func @log2_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
234224
// CHECK: %[[VAL_69:.*]] = subf %[[U]], %[[CST_ONE]] : f32
235225
// CHECK: %[[VAL_70:.*]] = divf %[[LOG_U]], %[[VAL_69]] : f32
236226
// CHECK: %[[LOG_LARGE:.*]] = mulf %[[X]], %[[VAL_70]] : f32
237-
// CHECK: %[[VAL_72:.*]] = llvm.or %[[U_SMALL]], %[[U_INF]] : i1
227+
// CHECK: %[[VAL_72:.*]] = or %[[U_SMALL]], %[[U_INF]] : i1
238228
// CHECK: %[[APPROX:.*]] = select %[[VAL_72]], %[[X]], %[[LOG_LARGE]] : f32
239229
// CHECK: return %[[APPROX]] : f32
240230
// CHECK: }
@@ -246,8 +236,7 @@ func @log1p_scalar(%arg0: f32) -> f32 {
246236
// CHECK-LABEL: func @log1p_vector(
247237
// CHECK-SAME: %[[VAL_0:.*]]: vector<8xf32>) -> vector<8xf32> {
248238
// CHECK: %[[CST_ONE:.*]] = constant dense<1.000000e+00> : vector<8xf32>
249-
// CHECK-COUNT-6: vector.broadcast
250-
// CHECK-COUNT-5: select
239+
// CHECK-COUNT-6: select
251240
// CHECK: %[[VAL_79:.*]] = select
252241
// CHECK: return %[[VAL_79]] : vector<8xf32>
253242
// CHECK: }

0 commit comments

Comments
 (0)