llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

update AtenClampOp in torch-to-tosa to handle fp inputs (#2516) 

As titled.

---------

Co-authored-by: Ze Zhang <ze.zhang@getcruise.com>
pull/2517/head
Ze Zhang 2023-10-17 14:49:47 -07:00 committed by GitHub
parent 14a4da923b
commit 4279b750da
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 46 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
lib/Conversion/TorchToTosa/TorchToTosa.cpp
View File

@ -3984,6 +3984,23 @@ 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 (selfType.getElementType().isa<mlir::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(
@ -3993,10 +4010,11 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewrite(
return rewriter.notifyMatchFailure(
op, "unimplemented: value `int_max` should be a torch constant int");
IntegerAttr min_int = rewriter.getI64IntegerAttr(int_min);
IntegerAttr max_int = rewriter.getI64IntegerAttr(int_max);
FloatAttr min_fp = rewriter.getF32FloatAttr(float(int_min));
FloatAttr max_fp = rewriter.getF32FloatAttr(float(int_max));
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));
}
auto outType = getTypeConverter()->convertType(op.getType());
rewriter.replaceOpWithNewOp<tosa::ClampOp>(op, outType, adaptor.getSelf(),

17
test/Conversion/TorchToTosa/basic.mlir
View File

@ -1072,6 +1072,23 @@ func.func @torch.aten.clamp(%arg0: !torch.vtensor<[1,1,128,128],si64>) -> !torch
return %0 : !torch.vtensor<[1,1,128,128],si64>
}
// -----
// CHECK-LABEL: func.func @torch.aten.clamp.float(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[1,1,128,128],f32>) -> !torch.vtensor<[1,1,128,128],f32> {
// CHECK: %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[1,1,128,128],f32> -> tensor<1x1x128x128xf32>
// CHECK: %[[VAL_2:.*]] = torch.constant.float 3.123400e+00
// CHECK: %[[VAL_3:.*]] = torch.constant.float 6.432100e+00
// CHECK: %[[VAL_4:.*]] = tosa.clamp %[[VAL_1]] {max_fp = 6.432100e+00 : f32, max_int = 6 : i64, min_fp = 3.123400e+00 : f32, min_int = 3 : i64} : (tensor<1x1x128x128xf32>) -> tensor<1x1x128x128xf32>
// CHECK: %[[VAL_5:.*]] = torch_c.from_builtin_tensor %[[VAL_4]] : tensor<1x1x128x128xf32> -> !torch.vtensor<[1,1,128,128],f32>
// CHECK: return %[[VAL_5]] : !torch.vtensor<[1,1,128,128],f32>
// CHECK: }
func.func @torch.aten.clamp.float(%arg0: !torch.vtensor<[1,1,128,128],f32>) -> !torch.vtensor<[1,1,128,128],f32> {
%fp_min = torch.constant.float 3.123400e+00
%fp_max = torch.constant.float 6.432100e+00
%0 = torch.aten.clamp %arg0, %fp_min, %fp_max : !torch.vtensor<[1,1,128,128],f32>, !torch.float, !torch.float -> !torch.vtensor<[1,1,128,128],f32>
return %0 : !torch.vtensor<[1,1,128,128],f32>
}
// -----
// CHECK-LABEL: func.func @torch.aten.masked_fill.Scalar(
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[1,12,128,128],f32>,