Skip to content

Commit

Permalink
Handle torch.none type in tosa.clamp op (#2739)
Browse files Browse the repository at this point in the history 
This PR updates the torch-to-tosa conversion with following changes:

- Support torch.none as min/max input argument for tosa.clamp op
- Support negative value as start index for tosa.slice op
- Add tosa.logical_or lowering support

e2e test:
python -m e2e_testing.main --config=tosa

LIT tests:
cmake --build build --target tools/torch-mlir/all

---------

Co-authored-by: Ze Zhang <ze.zhang@getcruise.com>
  • Loading branch information
@zezhang
zezhang and Ze Zhang authored Jan 11, 2024
1 parent 47ffc90 commit 670a99a
Show file tree
Hide file tree
Showing 3 changed files with 129 additions and 32 deletions.
78 changes: 46 additions & 32 deletions lib/Conversion/TorchToTosa/TorchToTosa.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Traits.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/DialectConversion.h"
#include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
Expand Down Expand Up @@ -3336,9 +3337,11 @@ LogicalResult ConvertAtenOp<AtenSliceTensorOp>::matchAndRewrite(
if (!matchPattern(op.getStart(), m_TorchConstantInt(&start)))
return rewriter.notifyMatchFailure(op, "start must be a Scalar constant");

if (start < 0)
return rewriter.notifyMatchFailure(op, "Currently unsupported: start < 0");

if (start < 0) {
start = toPositiveDim(start, selfType.getShape()[dim]);
if (!isValidDim(start, selfType.getShape()[dim]))
return rewriter.notifyMatchFailure(op, "start is not a valid index");
}
start = std::min(selfType.getShape()[dim], start);

int64_t end;
Expand Down Expand Up @@ -3984,36 +3987,46 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewrite(
return rewriter.notifyMatchFailure(
op, "only tensor types input are currently supported");

IntegerAttr min_int, max_int;
FloatAttr min_fp, max_fp;
if (op.getMin().getType().isa<Torch::FloatType>()) {
double fp_min, fp_max;
if (!matchPattern(op.getMin(), m_TorchConstantFloat(&fp_min)))
return rewriter.notifyMatchFailure(
op, "unimplemented: value `fp_min` should be a torch constant float");

if (!matchPattern(op.getMax(), m_TorchConstantFloat(&fp_max)))
return rewriter.notifyMatchFailure(
op, "unimplemented: value `fp_max` should be a torch constant float");

min_int = rewriter.getI64IntegerAttr(static_cast<int64_t>(fp_min));
max_int = rewriter.getI64IntegerAttr(static_cast<int64_t>(fp_max));
min_fp = rewriter.getF32FloatAttr(static_cast<float>(fp_min));
max_fp = rewriter.getF32FloatAttr(static_cast<float>(fp_max));
} else {
int64_t int_min, int_max;
if (!matchPattern(op.getMin(), m_TorchConstantInt(&int_min)))
return rewriter.notifyMatchFailure(
op, "unimplemented: value `int_min` should be a torch constant int");

if (!matchPattern(op.getMax(), m_TorchConstantInt(&int_max)))
return rewriter.notifyMatchFailure(
op, "unimplemented: value `int_max` should be a torch constant int");
IntegerAttr min_int =
rewriter.getI64IntegerAttr(std::numeric_limits<int64_t>::min());
IntegerAttr max_int =
rewriter.getI64IntegerAttr(std::numeric_limits<int64_t>::max());
FloatAttr min_fp =
rewriter.getF32FloatAttr(std::numeric_limits<float>::lowest());
FloatAttr max_fp =
rewriter.getF32FloatAttr(std::numeric_limits<float>::max());

auto getValAttr = [&](Value operand, IntegerAttr &intAttr,
FloatAttr &fpAttr) -> LogicalResult {
double valFloat;
int64_t valInt;
if (matchPattern(operand, m_TorchConstantFloat(&valFloat))) {
intAttr = rewriter.getI64IntegerAttr(static_cast<int64_t>(valFloat));
fpAttr = rewriter.getF32FloatAttr(static_cast<float>(valFloat));
} else if (matchPattern(operand, m_TorchConstantInt(&valInt))) {
intAttr = rewriter.getI64IntegerAttr(valInt);
fpAttr = rewriter.getF32FloatAttr(static_cast<float>(valInt));
} else {
return failure();
}
return success();
};

min_int = rewriter.getI64IntegerAttr(int_min);
max_int = rewriter.getI64IntegerAttr(int_max);
min_fp = rewriter.getF32FloatAttr(static_cast<float>(int_min));
max_fp = rewriter.getF32FloatAttr(static_cast<float>(int_max));
LogicalResult minAttrResult = getValAttr(op.getMin(), min_int, min_fp);
LogicalResult maxAttrResult = getValAttr(op.getMax(), max_int, max_fp);
if (failed(minAttrResult) && failed(maxAttrResult)) {
return rewriter.notifyMatchFailure(
op, "either `min` or `max` should be a torch constant");
}
if (failed(minAttrResult) &&
succeeded(checkNotNone(rewriter, op, op.getMin()))) {
return rewriter.notifyMatchFailure(op,
"min attr should be a torch constant");
}
if (failed(maxAttrResult) &&
succeeded(checkNotNone(rewriter, op, op.getMax()))) {
return rewriter.notifyMatchFailure(op,
"max attr should be a torch constant");
}

auto outType = getTypeConverter()->convertType(op.getType());
Expand Down Expand Up @@ -5025,6 +5038,7 @@ class ConvertTorchToTosa : public ConvertTorchToTosaBase<ConvertTorchToTosa> {
patterns.add<ConvertAtenBinaryOp<AtenOp, TosaOp>>(typeConverter, context);
INSERT_BINARY_PATTERN(AtenMaximumOp, tosa::MaximumOp)
INSERT_BINARY_PATTERN(AtenMinimumOp, tosa::MinimumOp)
INSERT_BINARY_PATTERN(AtenLogicalOrOp, tosa::LogicalOrOp)
#undef INSERT_BINARY_PATTERN

#define INSERT_BINARY_ADDSUB_PATTERN(AtenOp, TosaOp) \
Expand Down
12 changes: 12 additions & 0 deletions projects/pt1/e2e_testing/xfail_sets.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1035,6 +1035,15 @@
"ElementwiseAddScalar_TensorLiteralInt32_Module_basic",
"ElementwiseAtenDivIntScalarModule_basic",
"ElementwiseAtenIsinfOpModule_basic",
"ElementwiseAtenLogicalOrOpBrodcastModule_basic",
"ElementwiseAtenLogicalOrOpDiffArgs1Module_basic",
"ElementwiseAtenLogicalOrOpDiffArgs2Module_basic",
"ElementwiseAtenLogicalOrOpDiffArgs3Module_basic",
"ElementwiseAtenLogicalOrOpModule_basic",
"ElementwiseAtenLogicalOrOpNegativeModule_basic",
"ElementwiseAtenLogicalOrOpPromoteBroadcastStaticShapeModule_basic",
"ElementwiseAtenLogicalOrOpRandomFloatModule_basic",
"ElementwiseAtenLogicalOrOpRandomModule_basic",
"ElementwiseAtenWhereSelfModule_basic",
"ElementwiseBinaryModule_basic",
"ElementwiseBinaryStaticShapeModule_basic",
Expand All @@ -1047,6 +1056,9 @@
"ElementwiseBitwiseXorModule_basic",
"ElementwiseBitwiseXorStaticShapeModule_basic",
"ElementwiseCeilModule_basic",
"ElementwiseClampMaxModule_basic",
"ElementwiseClampMinModule_basic",
"ElementwiseClampModule_basic",
"ElementwiseCloneChannelsLastMemoryFormatModule_basic",
"ElementwiseCloneContiguousModule_basic",
"ElementwiseCloneModule_basic",
Expand Down
71 changes: 71 additions & 0 deletions test/Conversion/TorchToTosa/basic.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,22 @@ func.func @torch.aten.ne.Tensor$basic(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !

// -----

// CHECK-LABEL: func.func @torch.aten.logical_or$basic(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[?,?],i1>,
// CHECK-SAME: %[[VAL_1:.*]]: !torch.vtensor<[?,?],i1>) -> !torch.vtensor<[?,?],i1> {
// CHECK: %[[VAL_2:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[?,?],i1> -> tensor<?x?xi1>
// CHECK: %[[VAL_3:.*]] = torch_c.to_builtin_tensor %[[VAL_1]] : !torch.vtensor<[?,?],i1> -> tensor<?x?xi1>
// CHECK: %[[VAL_4:.*]] = tosa.logical_or %[[VAL_2]], %[[VAL_3]] : (tensor<?x?xi1>, tensor<?x?xi1>) -> tensor<?x?xi1>
// CHECK: %[[VAL_5:.*]] = torch_c.from_builtin_tensor %[[VAL_4]] : tensor<?x?xi1> -> !torch.vtensor<[?,?],i1>
// CHECK: return %[[VAL_5]] : !torch.vtensor<[?,?],i1>
// CHECK: }
func.func @torch.aten.logical_or$basic(%arg0: !torch.vtensor<[?,?],i1>, %arg1: !torch.vtensor<[?,?],i1>) -> !torch.vtensor<[?,?],i1> {
%0 = torch.aten.logical_or %arg0, %arg1 : !torch.vtensor<[?,?],i1>, !torch.vtensor<[?,?],i1> -> !torch.vtensor<[?,?],i1>
return %0 : !torch.vtensor<[?,?],i1>
}

// -----

// CHECK-LABEL: func.func @forward(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[3,4,2],f32>) -> !torch.vtensor<[3,2,4],f32> {
// CHECK: %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[3,4,2],f32> -> tensor<3x4x2xf32>
Expand Down Expand Up @@ -1055,6 +1071,61 @@ func.func @torch.aten.Scalar$basic(%arg0: !torch.vtensor<[1,1,128,128],si64>) ->
return %0 : !torch.vtensor<[1,1,128,128],si64>
}

// -----
// CHECK-LABEL: func.func @torch.aten.slice.negative_start(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[4,65,256],f32>) -> !torch.vtensor<[4,16,256],f32> {
// CHECK: %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[4,65,256],f32> -> tensor<4x65x256xf32>
// CHECK: %[[VAL_2:.*]] = torch.constant.int 0
// CHECK: %[[VAL_3:.*]] = torch.constant.int 1
// CHECK: %[[VAL_4:.*]] = torch.constant.int 100
// CHECK: %[[VAL_5:.*]] = torch.constant.int -16
// CHECK: %[[VAL_4:.*]] = tosa.slice %[[VAL_1]] {size = array<i64: 4, 16, 256>, start = array<i64: 0, 49, 0>} : (tensor<4x65x256xf32>) -> tensor<4x16x256xf32>
// CHECK: %[[VAL_5:.*]] = torch_c.from_builtin_tensor %[[VAL_4]] : tensor<4x16x256xf32> -> !torch.vtensor<[4,16,256],f32>
// CHECK: return %[[VAL_5]] : !torch.vtensor<[4,16,256],f32>
// CHECK: }
func.func @torch.aten.slice.negative_start(%arg0: !torch.vtensor<[4,65,256],f32>) -> !torch.vtensor<[4,16,256],f32> {
%int0 = torch.constant.int 0
%int1 = torch.constant.int 1
%int100 = torch.constant.int 100
%int-16 = torch.constant.int -16
%0 = torch.aten.slice.Tensor %arg0, %int1, %int-16, %int100, %int1 : !torch.vtensor<[4,65,256],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[4,16,256],f32>
return %0 : !torch.vtensor<[4,16,256],f32>
}

// -----
// CHECK-LABEL: func.func @torch.aten.clamp.min_none(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[1,1,128,128],si64>) -> !torch.vtensor<[1,1,128,128],si64> {
// CHECK: %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[1,1,128,128],si64> -> tensor<1x1x128x128xi64>
// CHECK: %[[VAL_2:.*]] = torch.constant.int 0
// CHECK: %[[VAL_3:.*]] = torch.constant.none
// CHECK: %[[VAL_4:.*]] = tosa.clamp %[[VAL_1]] {max_fp = 0.000000e+00 : f32, max_int = 0 : i64, min_fp = -3.40282347E+38 : f32, min_int = -9223372036854775808 : i64} : (tensor<1x1x128x128xi64>) -> tensor<1x1x128x128xi64>
// CHECK: %[[VAL_5:.*]] = torch_c.from_builtin_tensor %[[VAL_4]] : tensor<1x1x128x128xi64> -> !torch.vtensor<[1,1,128,128],si64>
// CHECK: return %[[VAL_5]] : !torch.vtensor<[1,1,128,128],si64>
// CHECK: }
func.func @torch.aten.clamp.min_none(%arg0: !torch.vtensor<[1,1,128,128],si64>) -> !torch.vtensor<[1,1,128,128],si64> {
%int0 = torch.constant.int 0
%none = torch.constant.none
%0 = torch.aten.clamp %arg0, %none, %int0 : !torch.vtensor<[1,1,128,128],si64>, !torch.none, !torch.int -> !torch.vtensor<[1,1,128,128],si64>
return %0 : !torch.vtensor<[1,1,128,128],si64>
}

// -----
// CHECK-LABEL: func.func @torch.aten.clamp.max_none(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[1,1,128,128],si64>) -> !torch.vtensor<[1,1,128,128],si64> {
// CHECK: %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[1,1,128,128],si64> -> tensor<1x1x128x128xi64>
// CHECK: %[[VAL_2:.*]] = torch.constant.int 0
// CHECK: %[[VAL_3:.*]] = torch.constant.none
// CHECK: %[[VAL_4:.*]] = tosa.clamp %[[VAL_1]] {max_fp = 3.40282347E+38 : f32, max_int = 9223372036854775807 : i64, min_fp = 0.000000e+00 : f32, min_int = 0 : i64} : (tensor<1x1x128x128xi64>) -> tensor<1x1x128x128xi64>
// CHECK: %[[VAL_5:.*]] = torch_c.from_builtin_tensor %[[VAL_4]] : tensor<1x1x128x128xi64> -> !torch.vtensor<[1,1,128,128],si64>
// CHECK: return %[[VAL_5]] : !torch.vtensor<[1,1,128,128],si64>
// CHECK: }
func.func @torch.aten.clamp.max_none(%arg0: !torch.vtensor<[1,1,128,128],si64>) -> !torch.vtensor<[1,1,128,128],si64> {
%int0 = torch.constant.int 0
%none = torch.constant.none
%0 = torch.aten.clamp %arg0, %int0, %none : !torch.vtensor<[1,1,128,128],si64>, !torch.int, !torch.none -> !torch.vtensor<[1,1,128,128],si64>
return %0 : !torch.vtensor<[1,1,128,128],si64>
}

// -----
// CHECK-LABEL: func.func @torch.aten.clamp(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[1,1,128,128],si64>) -> !torch.vtensor<[1,1,128,128],si64> {
Expand Down

0 comments on commit 670a99a

Please sign in to comment.