Skip to content

Commit 1f35b31

Browse files
oowekyalaoowekyala
committed
Cleanup handling of max/min in reduction
Since cinm is a funnel from/to linalg it should be as expressive. Previously it was discarding info about signedness of integers and NaN behavior with floats.
1 parent bc37ed7 commit 1f35b31

10 files changed

Lines changed: 155 additions & 171 deletions

File tree

‎include/cinm-mlir/Dialect/Cinm/IR/CinmAttributesForOps.td‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,16 @@ def Cinm_ReduceMethodAttr : I64EnumAttr<
2626
[
2727
I64EnumAttrCase<"ADD", 0, "add">,
2828
I64EnumAttrCase<"MUL", 1, "mul">,
29-
I64EnumAttrCase<"MAX", 2, "max">,
30-
I64EnumAttrCase<"MIN", 3, "min">,
29+
30+
I64EnumAttrCase<"MAXSI", 2, "maxsi">,
31+
I64EnumAttrCase<"MAXUI", 3, "maxui">,
32+
I64EnumAttrCase<"MAXIMUMF", 4, "maximumf">,
33+
I64EnumAttrCase<"MAXNUMF", 5, "maxnumf">,
34+
35+
I64EnumAttrCase<"MINSI", 6, "minsi">,
36+
I64EnumAttrCase<"MINUI", 7, "minui">,
37+
I64EnumAttrCase<"MINIMUMF", 8, "minimumf">,
38+
I64EnumAttrCase<"MINNUMF", 9, "minnumf">,
3139
]> {
3240
let cppNamespace = "::mlir::cinm";
3341
}

‎include/cinm-mlir/Dialect/Cinm/IR/CinmOps.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ namespace mlir::cinm {
3030

3131
Type inferGemmReturnType(Type lhsType, Type rhsType);
3232

33+
arith::AtomicRMWKind getArithConstant(ReduceMethod r, Type ty);
3334
cinm::ComputeBlockOp getEnclosingComputeBlock(Operation *op);
3435
cinm::CinmAcceleratorAttrInterface getEnclosingAccelerator(Operation *op);
3536
template <class T> T getEnclosingAcceleratorAs(Operation *op) {

‎lib/Conversion/CinmToCnm/CinmToCnm.cpp‎

Lines changed: 5 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1004,8 +1004,6 @@ struct ConvertCinmReduceToCnm : public OpConversionPattern<cinm::ReduceOp> {
10041004
builder, op.getResult().getType(),
10051005
builder.getZeroAttr(op.getResult().getType()));
10061006

1007-
const bool isFloatOp =
1008-
isa<FloatType>(cast<ShapedType>(op.getType()).getElementType());
10091007
SmallVector<int64_t> redDim = {static_cast<int64_t>(op.getDimension())};
10101008

10111009
llvm::SmallVector<Value, 1> newResults;
@@ -1019,45 +1017,11 @@ struct ConvertCinmReduceToCnm : public OpConversionPattern<cinm::ReduceOp> {
10191017
builder, inputs, outputs, ArrayRef<int64_t>{0},
10201018
[&](OpBuilder &builder, Location loc,
10211019
ValueRange inputs) -> void {
1022-
Value result;
1023-
switch (op.getMethod()) {
1024-
case mlir::cinm::ReduceMethod::ADD: {
1025-
if (isFloatOp) {
1026-
result = arith::AddFOp::create(builder, loc, inputs[0],
1027-
inputs[1]);
1028-
} else {
1029-
result = arith::AddIOp::create(builder, loc, inputs[0],
1030-
inputs[1]);
1031-
}
1032-
} break;
1033-
case mlir::cinm::ReduceMethod::MUL: {
1034-
if (isFloatOp) {
1035-
result = arith::MulFOp::create(builder, loc, inputs[0],
1036-
inputs[1]);
1037-
} else {
1038-
result = arith::MulIOp::create(builder, loc, inputs[0],
1039-
inputs[1]);
1040-
}
1041-
} break;
1042-
case mlir::cinm::ReduceMethod::MAX: {
1043-
if (isFloatOp) {
1044-
result = arith::MaximumFOp::create(
1045-
builder, loc, inputs[0], inputs[1]);
1046-
} else {
1047-
result = arith::MaxSIOp::create(builder, loc, inputs[0],
1048-
inputs[1]);
1049-
}
1050-
} break;
1051-
case mlir::cinm::ReduceMethod::MIN: {
1052-
if (isFloatOp) {
1053-
result = arith::MinimumFOp::create(
1054-
builder, loc, inputs[0], inputs[1]);
1055-
} else {
1056-
result = arith::MinSIOp::create(builder, loc, inputs[0],
1057-
inputs[1]);
1058-
}
1059-
} break;
1060-
}
1020+
arith::AtomicRMWKind arithMethod = cinm::getArithConstant(
1021+
op.getMethod(),
1022+
op.getInput().getType().getElementType());
1023+
Value result = arith::getReductionOp(
1024+
arithMethod, builder, loc, inputs[0], inputs[1]);
10611025
linalg::YieldOp::create(builder, loc, result);
10621026
});
10631027
})

‎lib/Conversion/CinmToLinalg/CinmToLinalg.cpp‎

Lines changed: 12 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -52,59 +52,14 @@ static Value buildZero(OpBuilder &b, Location loc, Type elemType) {
5252

5353
static Value buildReduceIdentity(OpBuilder &b, Location loc,
5454
ReduceMethod method, Type elemType) {
55-
if (isa<FloatType>(elemType)) {
56-
const auto &sem = cast<FloatType>(elemType).getFloatSemantics();
57-
switch (method) {
58-
case ReduceMethod::ADD:
59-
return b.create<arith::ConstantOp>(
60-
loc, FloatAttr::get(elemType, APFloat::getZero(sem)));
61-
case ReduceMethod::MUL:
62-
return b.create<arith::ConstantOp>(
63-
loc, FloatAttr::get(elemType, APFloat(sem, 1)));
64-
case ReduceMethod::MAX:
65-
return b.create<arith::ConstantOp>(
66-
loc,
67-
FloatAttr::get(elemType, APFloat::getInf(sem, /*negative=*/true)));
68-
case ReduceMethod::MIN:
69-
return b.create<arith::ConstantOp>(
70-
loc,
71-
FloatAttr::get(elemType, APFloat::getInf(sem, /*negative=*/false)));
72-
}
73-
}
74-
unsigned width = cast<IntegerType>(elemType).getWidth();
75-
switch (method) {
76-
case ReduceMethod::ADD:
77-
return b.create<arith::ConstantIntOp>(loc, elemType, 0);
78-
case ReduceMethod::MUL:
79-
return b.create<arith::ConstantIntOp>(loc, elemType, 1);
80-
case ReduceMethod::MAX:
81-
return b.create<arith::ConstantOp>(
82-
loc, IntegerAttr::get(elemType, APInt::getSignedMinValue(width)));
83-
case ReduceMethod::MIN:
84-
return b.create<arith::ConstantOp>(
85-
loc, IntegerAttr::get(elemType, APInt::getSignedMaxValue(width)));
86-
}
87-
llvm_unreachable("unknown ReduceMethod");
55+
auto arithConst = cinm::getArithConstant(method, elemType);
56+
return arith::getIdentityValue(arithConst, elemType, b, loc);
8857
}
8958

9059
static Value emitReduceCombine(OpBuilder &b, Location loc, ReduceMethod method,
9160
Value elem, Value acc, Type elemType) {
92-
bool isFloat = isa<FloatType>(elemType);
93-
switch (method) {
94-
case ReduceMethod::ADD:
95-
return isFloat ? b.create<arith::AddFOp>(loc, elem, acc).getResult()
96-
: b.create<arith::AddIOp>(loc, elem, acc).getResult();
97-
case ReduceMethod::MUL:
98-
return isFloat ? b.create<arith::MulFOp>(loc, elem, acc).getResult()
99-
: b.create<arith::MulIOp>(loc, elem, acc).getResult();
100-
case ReduceMethod::MAX:
101-
return isFloat ? b.create<arith::MaxNumFOp>(loc, elem, acc).getResult()
102-
: b.create<arith::MaxSIOp>(loc, elem, acc).getResult();
103-
case ReduceMethod::MIN:
104-
return isFloat ? b.create<arith::MinNumFOp>(loc, elem, acc).getResult()
105-
: b.create<arith::MinSIOp>(loc, elem, acc).getResult();
106-
}
107-
llvm_unreachable("unknown ReduceMethod");
61+
auto arithConst = cinm::getArithConstant(method, elemType);
62+
return arith::getReductionOp(arithConst, b, loc, elem, acc);
10863
}
10964

11065
// Build the outs init for a gemm-like op, in priority order:
@@ -409,7 +364,8 @@ struct ConvertGemvToLinalg : public OpConversionPattern<cinm::GemvOp> {
409364

410365
auto loc = op.getLoc();
411366
auto resultTy = cast<RankedTensorType>(op.getResult().getType());
412-
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(), adaptor.getBias(), resultTy);
367+
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(),
368+
adaptor.getBias(), resultTy);
413369

414370
Value result = rewriter
415371
.create<linalg::MatvecOp>(
@@ -437,7 +393,8 @@ struct ConvertGemmToLinalg : public OpConversionPattern<cinm::GemmOp> {
437393

438394
auto loc = op.getLoc();
439395
auto resultTy = cast<RankedTensorType>(op.getResult().getType());
440-
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(), adaptor.getBias(), resultTy);
396+
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(),
397+
adaptor.getBias(), resultTy);
441398

442399
Value result = rewriter
443400
.create<linalg::MatmulOp>(
@@ -466,7 +423,8 @@ struct ConvertBatchGemmToLinalg
466423

467424
auto loc = op.getLoc();
468425
auto resultTy = cast<RankedTensorType>(op.getResult().getType());
469-
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(), adaptor.getBias(), resultTy);
426+
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(),
427+
adaptor.getBias(), resultTy);
470428

471429
Value result = rewriter
472430
.create<linalg::BatchMatmulOp>(
@@ -513,7 +471,8 @@ struct ConvertBatchGemvToLinalg
513471
linalg::IteratorTypeAttr::get(ctx, utils::IteratorType::reduction),
514472
};
515473

516-
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(), adaptor.getBias(), resultTy);
474+
Value init = buildGemmInit(rewriter, loc, adaptor.getOut(),
475+
adaptor.getBias(), resultTy);
517476
auto generic = rewriter.create<linalg::GenericOp>(
518477
loc, TypeRange{resultTy},
519478
ValueRange{adaptor.getLhs(), adaptor.getRhs()}, ValueRange{init},

‎lib/Conversion/LinalgToCinm/LinalgToCinm.cpp‎

Lines changed: 39 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1229,28 +1229,21 @@ template <> struct LinalgToCinmOpBuilder<linalg::ReduceOp> {
12291229
}
12301230
}
12311231

1232-
if (llvm::dyn_cast<arith::AddFOp>(reductionOp) ||
1233-
llvm::dyn_cast<arith::AddIOp>(reductionOp)) {
1234-
return cinm::ReduceMethod::ADD;
1235-
}
1236-
1237-
if (llvm::dyn_cast<arith::MulFOp>(reductionOp) ||
1238-
llvm::dyn_cast<arith::MulIOp>(reductionOp)) {
1239-
return cinm::ReduceMethod::MUL;
1240-
}
1241-
1242-
if (llvm::dyn_cast<arith::MinimumFOp>(reductionOp) ||
1243-
llvm::dyn_cast<arith::MinSIOp>(reductionOp)) {
1244-
return cinm::ReduceMethod::MIN;
1245-
}
1246-
1247-
if (llvm::dyn_cast<arith::MaximumFOp>(reductionOp) ||
1248-
llvm::dyn_cast<arith::MaxSIOp>(reductionOp)) {
1249-
return cinm::ReduceMethod::MAX;
1250-
}
1251-
1252-
// unsupported reduction method
1253-
return failure();
1232+
return llvm::TypeSwitch<Operation *, FailureOr<cinm::ReduceMethod>>(
1233+
reductionOp)
1234+
.Case([](arith::AddFOp) { return cinm::ReduceMethod::ADD; })
1235+
.Case([](arith::AddIOp) { return cinm::ReduceMethod::ADD; })
1236+
.Case([](arith::MulFOp) { return cinm::ReduceMethod::MUL; })
1237+
.Case([](arith::MulIOp) { return cinm::ReduceMethod::MUL; })
1238+
.Case([](arith::MaxSIOp) { return cinm::ReduceMethod::MAXSI; })
1239+
.Case([](arith::MaxUIOp) { return cinm::ReduceMethod::MAXUI; })
1240+
.Case([](arith::MaxNumFOp) { return cinm::ReduceMethod::MAXNUMF; })
1241+
.Case([](arith::MaximumFOp) { return cinm::ReduceMethod::MAXIMUMF; })
1242+
.Case([](arith::MinSIOp) { return cinm::ReduceMethod::MINSI; })
1243+
.Case([](arith::MinUIOp) { return cinm::ReduceMethod::MINUI; })
1244+
.Case([](arith::MinNumFOp) { return cinm::ReduceMethod::MINNUMF; })
1245+
.Case([](arith::MinimumFOp) { return cinm::ReduceMethod::MINIMUMF; })
1246+
.Default([](Operation *) { return failure(); });
12541247
}
12551248

12561249
static FailureOr<cinm::ReduceOp> build(ConversionPatternRewriter &rewriter,
@@ -1557,25 +1550,45 @@ struct ConvertLinalgGenericOpToCinm : OpConversionPattern<linalg::GenericOp> {
15571550
srcOp->getResultTypes()[0],
15581551
cinm::ReduceMethod::MUL, operand);
15591552
})
1553+
.Case<arith::MinUIOp>([&](arith::MinUIOp) {
1554+
return cinm::ReduceOp::create(rewriter, loc,
1555+
srcOp->getResultTypes()[0],
1556+
cinm::ReduceMethod::MINUI, operand);
1557+
})
15601558
.Case<arith::MinSIOp>([&](arith::MinSIOp) {
15611559
return cinm::ReduceOp::create(rewriter, loc,
15621560
srcOp->getResultTypes()[0],
1563-
cinm::ReduceMethod::MIN, operand);
1561+
cinm::ReduceMethod::MINSI, operand);
15641562
})
15651563
.Case<arith::MinimumFOp>([&](arith::MinimumFOp) {
15661564
return cinm::ReduceOp::create(rewriter, loc,
15671565
srcOp->getResultTypes()[0],
1568-
cinm::ReduceMethod::MIN, operand);
1566+
cinm::ReduceMethod::MINIMUMF, operand);
1567+
})
1568+
.Case<arith::MinNumFOp>([&](arith::MinNumFOp) {
1569+
return cinm::ReduceOp::create(rewriter, loc,
1570+
srcOp->getResultTypes()[0],
1571+
cinm::ReduceMethod::MINNUMF, operand);
15691572
})
15701573
.Case<arith::MaxSIOp>([&](arith::MaxSIOp) {
15711574
return cinm::ReduceOp::create(rewriter, loc,
15721575
srcOp->getResultTypes()[0],
1573-
cinm::ReduceMethod::MAX, operand);
1576+
cinm::ReduceMethod::MAXSI, operand);
1577+
})
1578+
.Case<arith::MaxUIOp>([&](arith::MaxUIOp) {
1579+
return cinm::ReduceOp::create(rewriter, loc,
1580+
srcOp->getResultTypes()[0],
1581+
cinm::ReduceMethod::MAXUI, operand);
15741582
})
15751583
.Case<arith::MaximumFOp>([&](arith::MaximumFOp) {
15761584
return cinm::ReduceOp::create(rewriter, loc,
15771585
srcOp->getResultTypes()[0],
1578-
cinm::ReduceMethod::MAX, operand);
1586+
cinm::ReduceMethod::MAXIMUMF, operand);
1587+
})
1588+
.Case<arith::MaxNumFOp>([&](arith::MaxNumFOp) {
1589+
return cinm::ReduceOp::create(rewriter, loc,
1590+
srcOp->getResultTypes()[0],
1591+
cinm::ReduceMethod::MAXNUMF, operand);
15791592
})
15801593
.Default([&](Operation *) {
15811594
// unsupported elementwise operation

‎lib/Dialect/Cinm/IR/CinmOps.cpp‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1002,3 +1002,37 @@ void ReduceOp::getCanonicalizationPatterns(::mlir::RewritePatternSet &results,
10021002
::mlir::MLIRContext *context) {
10031003
results.insert<ReduceOpNormalizeDim>(context);
10041004
}
1005+
1006+
arith::AtomicRMWKind cinm::getArithConstant(ReduceMethod r, Type ty) {
1007+
switch (r) {
1008+
case mlir::cinm::ReduceMethod::ADD:
1009+
if (ty.isFloat()) {
1010+
return mlir::arith::AtomicRMWKind::addf;
1011+
} else {
1012+
return mlir::arith::AtomicRMWKind::addi;
1013+
}
1014+
case mlir::cinm::ReduceMethod::MUL:
1015+
if (ty.isFloat()) {
1016+
return mlir::arith::AtomicRMWKind::mulf;
1017+
} else {
1018+
return mlir::arith::AtomicRMWKind::muli;
1019+
}
1020+
case mlir::cinm::ReduceMethod::MAXSI:
1021+
return mlir::arith::AtomicRMWKind::maxs;
1022+
case mlir::cinm::ReduceMethod::MAXUI:
1023+
return mlir::arith::AtomicRMWKind::maxu;
1024+
case mlir::cinm::ReduceMethod::MAXIMUMF:
1025+
return mlir::arith::AtomicRMWKind::maximumf;
1026+
case mlir::cinm::ReduceMethod::MAXNUMF:
1027+
return mlir::arith::AtomicRMWKind::maxnumf;
1028+
1029+
case mlir::cinm::ReduceMethod::MINSI:
1030+
return mlir::arith::AtomicRMWKind::mins;
1031+
case mlir::cinm::ReduceMethod::MINUI:
1032+
return mlir::arith::AtomicRMWKind::minu;
1033+
case mlir::cinm::ReduceMethod::MINIMUMF:
1034+
return mlir::arith::AtomicRMWKind::minimumf;
1035+
case mlir::cinm::ReduceMethod::MINNUMF:
1036+
return mlir::arith::AtomicRMWKind::minnumf;
1037+
}
1038+
}

‎lib/Dialect/Cinm/IR/CinmTilingImplementations.cpp‎

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -204,35 +204,6 @@ void ActivateOp::getTilableDimSizes(SmallVectorImpl<int64_t> &dimSizes) {
204204
// convertToTiledOps implementations
205205
// ---------------------------------------------------------------------------
206206

207-
static arith::AtomicRMWKind getArithConstant(ReduceMethod r, Type ty) {
208-
switch (r) {
209-
case mlir::cinm::ReduceMethod::ADD:
210-
if (ty.isFloat()) {
211-
return mlir::arith::AtomicRMWKind::addf;
212-
} else {
213-
return mlir::arith::AtomicRMWKind::addi;
214-
}
215-
case mlir::cinm::ReduceMethod::MUL:
216-
if (ty.isFloat()) {
217-
return mlir::arith::AtomicRMWKind::mulf;
218-
} else {
219-
return mlir::arith::AtomicRMWKind::muli;
220-
}
221-
case mlir::cinm::ReduceMethod::MAX:
222-
if (ty.isFloat()) {
223-
return mlir::arith::AtomicRMWKind::maximumf;
224-
} else {
225-
return mlir::arith::AtomicRMWKind::maxu;
226-
}
227-
case mlir::cinm::ReduceMethod::MIN:
228-
if (ty.isFloat()) {
229-
return mlir::arith::AtomicRMWKind::minimumf;
230-
} else {
231-
return mlir::arith::AtomicRMWKind::minu;
232-
}
233-
}
234-
}
235-
236207
static TypedAttr getNeutralElement(ReduceMethod r, Type ty, OpBuilder &builder,
237208
Location loc) {
238209
return arith::getIdentityValueAttr(getArithConstant(r, ty), ty, builder, loc);

‎lib/Dialect/Cinm/Transforms/SoftmaxToCinmPass.cpp‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -30,25 +30,28 @@ struct SoftmaxToCinmPattern : OpConversionPattern<linalg::SoftmaxOp> {
3030
ConversionPatternRewriter &rewriter) const override {
3131
const auto loc = op.getLoc();
3232
const ShapedType inputType = op.getInput().getType();
33+
bool isFloat = inputType.getElementType().isFloat();
3334

34-
auto compute =
35-
rewriter.replaceOpWithNewOp<ComputeBlockOp>(op, adaptor.getOperands(), op.getResultTypes());
35+
auto compute = rewriter.replaceOpWithNewOp<ComputeBlockOp>(
36+
op, adaptor.getOperands(), op.getResultTypes());
3637
Value innerInput = compute.getBodyArguments()[0];
3738

3839
rewriter.setInsertionPointToEnd(&compute.getBody().emplaceBlock());
39-
const Value max = rewriter.create<cinm::ReduceOp>(loc, inputType.getElementType(),
40-
ReduceMethod::MAX, innerInput, 0);
41-
const Value t =
42-
rewriter
43-
.create<cinm::ElementwiseOp>(loc, ElementwiseKind::Sub, innerInput, max)
44-
.getResult();
40+
const Value max = rewriter.create<cinm::ReduceOp>(
41+
loc, inputType.getElementType(),
42+
isFloat ? ReduceMethod::MAXNUMF : ReduceMethod::MAXSI, innerInput, 0);
43+
const Value t = rewriter
44+
.create<cinm::ElementwiseOp>(loc, ElementwiseKind::Sub,
45+
innerInput, max)
46+
.getResult();
4547
const SmallVector<Type, 1> types{RankedTensorType::get(
4648
inputType.getShape(), inputType.getElementType())};
4749

4850
const Value e =
49-
cinm::ElementwiseOp::create(rewriter, loc, ElementwiseKind::Exp, t).getResult();
50-
const Value s = rewriter.create<cinm::ReduceOp>(loc, inputType.getElementType(),
51-
ReduceMethod::ADD, e, 0);
51+
cinm::ElementwiseOp::create(rewriter, loc, ElementwiseKind::Exp, t)
52+
.getResult();
53+
const Value s = rewriter.create<cinm::ReduceOp>(
54+
loc, inputType.getElementType(), ReduceMethod::ADD, e, 0);
5255
const Value result =
5356
rewriter.create<cinm::ElementwiseOp>(loc, ElementwiseKind::Div, e, s)
5457
.getResult();

0 commit comments

Comments
 (0)