Skip to content

Commit a79579b

Browse files
ahagro01ahagro01
authored
[TorchToTosa] cast to match result (#4568)
Fixes aten.clamp conversion when the input tensor type differs from the result type. Previously, the lowering created a tosa.cast when promotion was needed, but still passed the original input to tosa.clamp. This MR updates the conversion to use the casted value, ensuring cases like si64 input with f32 clamp bounds lower correctly. Also adds a regression test covering integer-to-float promotion for torch.aten.clamp, verifying that the generated TOSA IR casts the input to f32 before applying tosa.clamp.
1 parent eaaa9c2 commit a79579b

4 files changed

Lines changed: 73 additions & 2 deletions

File tree

‎lib/Conversion/TorchToTosa/TorchToTosa.cpp‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6001,10 +6001,25 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewriteImpl(
60016001
if (!selfType)
60026002
return rewriter.notifyMatchFailure(
60036003
op, "only tensor types input are currently supported");
6004+
auto selfTorchType = dyn_cast<Torch::BaseTensorType>(op.getSelf().getType());
6005+
if (selfTorchType && selfTorchType.getDtype().isUnsignedInteger()) {
6006+
return rewriter.notifyMatchFailure(
6007+
op, "unsigned integer clamp is not currently supported");
6008+
}
60046009

60056010
auto outType =
60066011
dyn_cast<TensorType>(getTypeConverter()->convertType(op.getType()));
6012+
if (!outType)
6013+
return rewriter.notifyMatchFailure(
6014+
op, "only tensor types output are currently supported");
60076015
auto outElemTy = outType.getElementType();
6016+
Value self = adaptor.getSelf();
6017+
if (selfType != outType) {
6018+
auto castedSelf = tosa::tosaCastTensorToType(rewriter, self, outType);
6019+
if (!castedSelf)
6020+
return rewriter.notifyMatchFailure(op, "failed to cast self");
6021+
self = *castedSelf;
6022+
}
60086023

60096024
std::optional<int64_t> minInt;
60106025
std::optional<double> minFloat;
@@ -6048,7 +6063,7 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewriteImpl(
60486063
}
60496064

60506065
rewriter.replaceOpWithNewOp<tosa::ClampOp>(
6051-
op, outType, adaptor.getSelf(), minIntAttr, maxIntAttr,
6066+
op, outType, self, minIntAttr, maxIntAttr,
60526067
/*nan_mode=*/
60536068
tosa::NanPropagationModeAttr::get(rewriter.getContext(),
60546069
tosa::NanPropagationMode::PROPAGATE));
@@ -6061,7 +6076,7 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewriteImpl(
60616076
}
60626077

60636078
rewriter.replaceOpWithNewOp<tosa::ClampOp>(
6064-
op, outType, adaptor.getSelf(), minFloatAttr, maxFloatAttr,
6079+
op, outType, self, minFloatAttr, maxFloatAttr,
60656080
/*nan_mode=*/
60666081
tosa::NanPropagationModeAttr::get(rewriter.getContext(),
60676082
tosa::NanPropagationMode::PROPAGATE));

‎projects/pt1/e2e_testing/xfail_sets.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
"InterpolateStaticModule_sizes_bilinear_no_align_corners",
2525
"InterpolateDynamicModule_scales_recompute_bilinear",
2626
"ElementwiseFloatTensorGtIntTensorModule_basic",
27+
"ElementwiseClampIntToFloatModule_basic",
2728
# TODO: The values are extremely close to the golden values, but the test fails because of strict rtol/atol.
2829
"AtenInstanceNormModuleFp16_basic",
2930
"AtenIntMM_basic",
@@ -513,6 +514,7 @@
513514
"ReflectionPad3dModuleRight_basic",
514515
"ReflectionPad3dModuleFront_basic",
515516
"ReflectionPad3dModuleBack_basic",
517+
"ElementwiseClampIntToFloatModule_basic",
516518
# error: argument must be a memref of f32, f64, i32, i64, i8, i1, c32, c64, but got 'memref<3x5xbf16>'
517519
"ElementwiseClampMaxModule_bfloat16",
518520
"ElementwiseClampMinModule_bfloat16",
@@ -548,6 +550,7 @@
548550
"AtenPolarFloatModule_basic",
549551
"DiagonalWithStaticShapeModule_basic",
550552
"EinsumStaticDiagonalDimensionModule_basic",
553+
"ElementwiseClampIntToFloatModule_basic",
551554
"ElementwiseRemainderScalarModule_Bool_NegativeDivisor_basic",
552555
"ElementwiseRemainderScalarModule_Float_NegativeDividend_basic",
553556
"ElementwiseRemainderScalarModule_Float_NegativeDivisor_basic",
@@ -2762,6 +2765,7 @@
27622765
"ElementwiseAtan2TensorIntStaticModule_basic",
27632766
"ElementwiseAtenFloorDivideScalarNegativeModule_basic",
27642767
"ElementwiseAtenFloorDivideTensorNegativeModule_basic",
2768+
"ElementwiseClampIntToFloatModule_basic",
27652769
"ElementwiseLog10IntModule_basic",
27662770
"ElementwiseLog2IntModule_basic",
27672771
"ElementwiseFminModule_basic",
@@ -4074,6 +4078,7 @@
40744078
"L1LossMeanReductionModule_basic",
40754079
"L1LossNoReductionModule_basic",
40764080
"L1LossSumReductionModule_basic",
4081+
"ElementwiseClampIntToFloatModule_basic",
40774082
"ElementwiseLogAddExpModule_basic",
40784083
"ElementwiseLogAddExpBroadcastModule_basic",
40794084
"ElementwiseLogAddExp2Module_basic",

‎projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1815,6 +1815,29 @@ def ElementwiseClampModule_basic(module, tu: TestUtils):
18151815
# ==============================================================================
18161816

18171817

1818+
class ElementwiseClampIntToFloatModule(torch.nn.Module):
1819+
def __init__(self):
1820+
super().__init__()
1821+
1822+
@export
1823+
@annotate_args(
1824+
[
1825+
None,
1826+
([-1, -1], torch.int64, True),
1827+
]
1828+
)
1829+
def forward(self, x):
1830+
return torch.clamp(x, min=-2.5, max=2.5)
1831+
1832+
1833+
@register_test_case(module_factory=lambda: ElementwiseClampIntToFloatModule())
1834+
def ElementwiseClampIntToFloatModule_basic(module, tu: TestUtils):
1835+
module.forward(tu.randint(3, 5, low=-10, high=10))
1836+
1837+
1838+
# ==============================================================================
1839+
1840+
18181841
class ElementwiseClampBFloat16Module(torch.nn.Module):
18191842
def __init__(self):
18201843
super().__init__()

‎test/Conversion/TorchToTosa/basic.mlir‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1614,6 +1614,34 @@ func.func @torch.aten.clamp.float(%arg0: !torch.vtensor<[1,1,128,128],f32>) -> !
16141614
return %0 : !torch.vtensor<[1,1,128,128],f32>
16151615
}
16161616

1617+
// -----
1618+
// CHECK-LABEL: func.func @torch.aten.clamp.promote_int_input_to_float_result(
1619+
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[1,2500],si64>) -> !torch.vtensor<[1,2500],f32> {
1620+
// CHECK: %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[1,2500],si64> -> tensor<1x2500xi64>
1621+
// CHECK: %[[VAL_2:.*]] = torch.constant.float 1.500000e+00
1622+
// CHECK: %[[VAL_3:.*]] = torch.constant.none
1623+
// CHECK: %[[VAL_4:.*]] = tosa.cast %[[VAL_1]] : (tensor<1x2500xi64>) -> tensor<1x2500xf32>
1624+
// CHECK: %[[VAL_5:.*]] = tosa.clamp %[[VAL_4]] {max_val = 3.40282347E+38 : f32, min_val = 1.500000e+00 : f32} : (tensor<1x2500xf32>) -> tensor<1x2500xf32>
1625+
// CHECK: %[[VAL_6:.*]] = torch_c.from_builtin_tensor %[[VAL_5]] : tensor<1x2500xf32> -> !torch.vtensor<[1,2500],f32>
1626+
// CHECK: return %[[VAL_6]] : !torch.vtensor<[1,2500],f32>
1627+
// CHECK: }
1628+
func.func @torch.aten.clamp.promote_int_input_to_float_result(%arg0: !torch.vtensor<[1,2500],si64>) -> !torch.vtensor<[1,2500],f32> {
1629+
%fp_min = torch.constant.float 1.500000e+00
1630+
%none = torch.constant.none
1631+
%0 = torch.aten.clamp %arg0, %fp_min, %none : !torch.vtensor<[1,2500],si64>, !torch.float, !torch.none -> !torch.vtensor<[1,2500],f32>
1632+
return %0 : !torch.vtensor<[1,2500],f32>
1633+
}
1634+
1635+
// -----
1636+
1637+
func.func @torch.aten.clamp.unsigned_input(%arg0: !torch.vtensor<[1,2500],ui8>) -> !torch.vtensor<[1,2500],ui8> {
1638+
%int1 = torch.constant.int 1
1639+
%none = torch.constant.none
1640+
// expected-error @+1 {{failed to legalize operation 'torch.aten.clamp' that was explicitly marked illegal}}
1641+
%0 = torch.aten.clamp %arg0, %int1, %none : !torch.vtensor<[1,2500],ui8>, !torch.int, !torch.none -> !torch.vtensor<[1,2500],ui8>
1642+
return %0 : !torch.vtensor<[1,2500],ui8>
1643+
}
1644+
16171645
// -----
16181646
// CHECK-LABEL: func.func @torch.aten.masked_fill.Scalar(
16191647
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[1,12,128,128],f32>,

0 commit comments

Comments
 (0)