Skip to content

Commit 14e9340

Browse files
sahas3sahas3
authored
[TorchToLinalg] Fix max/min lowering for bool (#4701)
This PR fixes `i1` legalization for max/min ops. `aten.max`, `aten.min`, `aten.max.dim`, and `aten.min.dim` on bool (`i1`) tensors produced incorrect results or lowering failures. `i1` is signless, so it fell through both `isUnsigned()` and `isSigned()` checks; and even where it didn't fall through, signed compare treats bit `1` as `-1`, so `maxsi([false, true])` returns `false`.
1 parent 0c9fc54 commit 14e9340

4 files changed

Lines changed: 102 additions & 35 deletions

File tree

‎include/torch-mlir/Dialect/Torch/Utils/Utils.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,11 @@ LogicalResult getPermutedType(BaseTensorType inType,
158158
SmallVector<int64_t> permuteDims,
159159
Type &permutedType);
160160

161+
// Returns true when an IntegerType requires unsigned comparison/arithmetic
162+
// semantics: explicitly unsigned types, and signless i1 (bool) where unsigned
163+
// comparison correctly ranks true (1) above false (0).
164+
bool useUnsignedIntegerSemantics(IntegerType intType);
165+
161166
} // namespace Torch
162167
} // namespace torch
163168
} // namespace mlir

‎lib/Conversion/TorchToLinalg/Reduction.cpp‎

Lines changed: 39 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -92,12 +92,12 @@ class ConvertAtenMinMaxDimOp : public OpConversionPattern<OpTy> {
9292
return rewriter.notifyMatchFailure(op, "dim is not a valid dim");
9393

9494
Type inElementType = inputType.getElementType();
95-
bool isUnsigned = false;
95+
bool useUnsigned = false;
9696
if (!isa<mlir::FloatType>(inElementType)) {
9797
if (isa<mlir::IntegerType>(inElementType)) {
98-
auto integerTy = dyn_cast<mlir::IntegerType>(
98+
auto torchIntTy = dyn_cast<mlir::IntegerType>(
9999
cast<BaseTensorType>(op.getSelf().getType()).getDtype());
100-
isUnsigned = integerTy.isUnsigned();
100+
useUnsigned = useUnsignedIntegerSemantics(torchIntTy);
101101
} else {
102102
return rewriter.notifyMatchFailure(
103103
op, opName + " to linalg.* requires Float or Integer "
@@ -129,15 +129,15 @@ class ConvertAtenMinMaxDimOp : public OpConversionPattern<OpTy> {
129129
inElementType,
130130
getFloatInf(cast<mlir::FloatType>(inElementType),
131131
/*Negative=*/isMax, this->allowNonFinites)));
132-
} else if (!isUnsigned) {
133-
auto width = cast<mlir::IntegerType>(inElementType).getWidth();
134-
auto init = isMax ? APSInt::getSignedMinValue(width)
135-
: APSInt::getSignedMaxValue(width);
136-
fillValue = arith::ConstantOp::create(
137-
rewriter, loc, rewriter.getIntegerAttr(inElementType, init));
138-
} else if (isUnsigned) {
132+
} else {
139133
auto width = cast<mlir::IntegerType>(inElementType).getWidth();
140-
auto init = isMax ? APInt::getMinValue(width) : APInt::getMaxValue(width);
134+
APInt init;
135+
if (useUnsigned) {
136+
init = isMax ? APInt::getMinValue(width) : APInt::getMaxValue(width);
137+
} else {
138+
init = isMax ? APSInt::getSignedMinValue(width)
139+
: APSInt::getSignedMaxValue(width);
140+
}
141141
fillValue = arith::ConstantOp::create(
142142
rewriter, loc, rewriter.getIntegerAttr(inElementType, init));
143143
}
@@ -198,19 +198,19 @@ class ConvertAtenMinMaxDimOp : public OpConversionPattern<OpTy> {
198198
} else {
199199
arith::CmpIPredicate predType;
200200
if (isMax) {
201-
predType = isUnsigned ? arith::CmpIPredicate::ugt
202-
: arith::CmpIPredicate::sgt;
203-
if (isUnsigned) {
201+
predType = useUnsigned ? arith::CmpIPredicate::ugt
202+
: arith::CmpIPredicate::sgt;
203+
if (useUnsigned) {
204204
resultVal = arith::MaxUIOp::create(rewriter, nestedLoc,
205205
newValue, oldValue);
206206
} else {
207207
resultVal = arith::MaxSIOp::create(rewriter, nestedLoc,
208208
newValue, oldValue);
209209
}
210210
} else {
211-
predType = isUnsigned ? arith::CmpIPredicate::ult
212-
: arith::CmpIPredicate::slt;
213-
if (isUnsigned) {
211+
predType = useUnsigned ? arith::CmpIPredicate::ult
212+
: arith::CmpIPredicate::slt;
213+
if (useUnsigned) {
214214
resultVal = arith::MinUIOp::create(rewriter, nestedLoc,
215215
newValue, oldValue);
216216
} else {
@@ -318,12 +318,15 @@ static Value createInitElementForReduceOp(OpBuilder &b, Location loc,
318318
getFloatInf(cast<mlir::FloatType>(elementType),
319319
/*Negative=*/true, allowNonFinites)));
320320
else if (isa<mlir::IntegerType>(elementType) &&
321-
elementType.getIntOrFloatBitWidth() != 8)
322-
return arith::ConstantOp::create(
323-
b, loc,
324-
b.getIntegerAttr(
325-
elementType,
326-
APSInt::getSignedMinValue(elementType.getIntOrFloatBitWidth())));
321+
elementType.getIntOrFloatBitWidth() != 8) {
322+
unsigned width = elementType.getIntOrFloatBitWidth();
323+
auto init =
324+
useUnsignedIntegerSemantics(cast<mlir::IntegerType>(elementType))
325+
? APInt::getMinValue(width)
326+
: APSInt::getSignedMinValue(width);
327+
return arith::ConstantOp::create(b, loc,
328+
b.getIntegerAttr(elementType, init));
329+
}
327330
}
328331

329332
if (isa<AtenMinOp>(op)) {
@@ -334,12 +337,15 @@ static Value createInitElementForReduceOp(OpBuilder &b, Location loc,
334337
getFloatInf(cast<mlir::FloatType>(elementType),
335338
/*Negative=*/false, allowNonFinites)));
336339
else if (isa<mlir::IntegerType>(elementType) &&
337-
elementType.getIntOrFloatBitWidth() != 8)
338-
return arith::ConstantOp::create(
339-
b, loc,
340-
b.getIntegerAttr(
341-
elementType,
342-
APSInt::getSignedMaxValue(elementType.getIntOrFloatBitWidth())));
340+
elementType.getIntOrFloatBitWidth() != 8) {
341+
unsigned width = elementType.getIntOrFloatBitWidth();
342+
auto init =
343+
useUnsignedIntegerSemantics(cast<mlir::IntegerType>(elementType))
344+
? APInt::getMaxValue(width)
345+
: APSInt::getSignedMaxValue(width);
346+
return arith::ConstantOp::create(b, loc,
347+
b.getIntegerAttr(elementType, init));
348+
}
343349
}
344350

345351
if (isa<AtenLinalgVectorNormOp>(op) || isa<AtenFrobeniusNormDimOp>(op) ||
@@ -388,10 +394,9 @@ static Value createLinalgPayloadForReduceOp(OpBuilder &b, Location loc,
388394
else if (isa<mlir::IntegerType>(resultElementType)) {
389395
IntegerType intType = dyn_cast<mlir::IntegerType>(
390396
cast<BaseTensorType>(max.getSelf().getType()).getDtype());
391-
if (intType.isUnsigned())
397+
if (useUnsignedIntegerSemantics(intType))
392398
return arith::MaxUIOp::create(b, loc, self, result);
393-
if (intType.isSigned())
394-
return arith::MaxSIOp::create(b, loc, self, result);
399+
return arith::MaxSIOp::create(b, loc, self, result);
395400
}
396401
} else if (auto min = dyn_cast<AtenMinOp>(op)) {
397402
Value self =
@@ -402,10 +407,9 @@ static Value createLinalgPayloadForReduceOp(OpBuilder &b, Location loc,
402407
else if (isa<mlir::IntegerType>(resultElementType)) {
403408
IntegerType intType = dyn_cast<mlir::IntegerType>(
404409
cast<BaseTensorType>(min.getSelf().getType()).getDtype());
405-
if (intType.isUnsigned())
410+
if (useUnsignedIntegerSemantics(intType))
406411
return arith::MinUIOp::create(b, loc, self, result);
407-
if (intType.isSigned())
408-
return arith::MinSIOp::create(b, loc, self, result);
412+
return arith::MinSIOp::create(b, loc, self, result);
409413
}
410414
} else if (isa<AtenNormScalarOp>(op)) {
411415
// This creates payload for only the first of the two linalg.generic ops.

‎lib/Dialect/Torch/Utils/Utils.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@ using namespace mlir;
1818
using namespace mlir::torch;
1919
using namespace mlir::torch::Torch;
2020

21+
bool Torch::useUnsignedIntegerSemantics(IntegerType intType) {
22+
return intType.isUnsigned() || intType.getWidth() == 1;
23+
}
24+
2125
int64_t Torch::toPositiveDim(int64_t dim, int64_t inputRank) {
2226
return dim >= 0 ? dim : dim + inputRank;
2327
}

‎test/Conversion/TorchToLinalg/basic.mlir‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1121,3 +1121,57 @@ func.func @torch.aten.cat$rank1_empty(%arg0: !torch.vtensor<[1,8,?,128],f16>, %a
11211121
%1 = torch.aten.cat %0, %int-2 : !torch.list<vtensor>, !torch.int -> !torch.vtensor<[1,8,?,128],f16>
11221122
return %1 : !torch.vtensor<[1,8,?,128],f16>
11231123
}
1124+
1125+
// -----
1126+
// CHECK-LABEL: func.func @torch.aten.max$bool
1127+
// CHECK: %[[INIT:.+]] = arith.constant false
1128+
// CHECK: linalg.fill ins(%[[INIT]] : i1)
1129+
// CHECK: linalg.generic
1130+
// CHECK: ^bb0(%[[IN:.+]]: i1, %[[OUT:.+]]: i1):
1131+
// CHECK-NEXT: %[[RES:.+]] = arith.maxui %[[IN]], %[[OUT]] : i1
1132+
func.func @torch.aten.max$bool(%arg0: !torch.vtensor<[4],i1>) -> !torch.vtensor<[],i1> {
1133+
%0 = torch.aten.max %arg0 : !torch.vtensor<[4],i1> -> !torch.vtensor<[],i1>
1134+
return %0 : !torch.vtensor<[],i1>
1135+
}
1136+
1137+
// -----
1138+
// CHECK-LABEL: func.func @torch.aten.min$bool
1139+
// CHECK: %[[INIT:.+]] = arith.constant true
1140+
// CHECK: linalg.fill ins(%[[INIT]] : i1)
1141+
// CHECK: linalg.generic
1142+
// CHECK: ^bb0(%[[IN:.+]]: i1, %[[OUT:.+]]: i1):
1143+
// CHECK-NEXT: %[[RES:.+]] = arith.minui %[[IN]], %[[OUT]] : i1
1144+
func.func @torch.aten.min$bool(%arg0: !torch.vtensor<[4],i1>) -> !torch.vtensor<[],i1> {
1145+
%0 = torch.aten.min %arg0 : !torch.vtensor<[4],i1> -> !torch.vtensor<[],i1>
1146+
return %0 : !torch.vtensor<[],i1>
1147+
}
1148+
1149+
// -----
1150+
// CHECK-LABEL: func.func @torch.aten.max.dim$bool
1151+
// CHECK: %[[INIT:.+]] = arith.constant false
1152+
// CHECK: linalg.fill ins(%[[INIT]] : i1)
1153+
// CHECK: linalg.generic
1154+
// CHECK: ^bb0(%[[IN:.+]]: i1, {{.*}}: i1, {{.*}}: i64):
1155+
// CHECK: arith.maxui %[[IN]], {{.*}} : i1
1156+
// CHECK: arith.cmpi ugt, {{.*}} : i1
1157+
func.func @torch.aten.max.dim$bool(%arg0: !torch.vtensor<[3,4],i1>) -> !torch.vtensor<[3],i1> {
1158+
%false = torch.constant.bool false
1159+
%int1 = torch.constant.int 1
1160+
%values, %indices = torch.aten.max.dim %arg0, %int1, %false : !torch.vtensor<[3,4],i1>, !torch.int, !torch.bool -> !torch.vtensor<[3],i1>, !torch.vtensor<[3],si64>
1161+
return %values : !torch.vtensor<[3],i1>
1162+
}
1163+
1164+
// -----
1165+
// CHECK-LABEL: func.func @torch.aten.min.dim$bool
1166+
// CHECK: %[[INIT:.+]] = arith.constant true
1167+
// CHECK: linalg.fill ins(%[[INIT]] : i1)
1168+
// CHECK: linalg.generic
1169+
// CHECK: ^bb0(%[[IN:.+]]: i1, {{.*}}: i1, {{.*}}: i64):
1170+
// CHECK: arith.minui %[[IN]], {{.*}} : i1
1171+
// CHECK: arith.cmpi ult, {{.*}} : i1
1172+
func.func @torch.aten.min.dim$bool(%arg0: !torch.vtensor<[3,4],i1>) -> !torch.vtensor<[3],i1> {
1173+
%false = torch.constant.bool false
1174+
%int1 = torch.constant.int 1
1175+
%values, %indices = torch.aten.min.dim %arg0, %int1, %false : !torch.vtensor<[3,4],i1>, !torch.int, !torch.bool -> !torch.vtensor<[3],i1>, !torch.vtensor<[3],si64>
1176+
return %values : !torch.vtensor<[3],i1>
1177+
}

0 commit comments

Comments
 (0)