[TorchToTosa] Fix i32 and i64 aten.clamp lowering (#4576)

Tosa clamp does not support i32 and i64 operands so lowering aten.clamp
directly to tosa.clamp results in invalid tosa for these values.

This PR updates the lowering for scalar aten.clamp `AtenClampOp` so i8
and i16 operands continue to use tosa.clamp, while i32 and i64 are
lowered as `y = min(max(x, min_value), max_value)`, in the same way as
the existing lowering of the tensor aten.clamp `AtenClampTensorOp`

Author: jackv1ck

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -5991,13 +5991,93 @@ LogicalResult ConvertAtenOp<AtenIscloseOp>::matchAndRewriteImpl(
   return success();
 }
 
+static LogicalResult
+rewriteClampAsMinimumMaximumOp(Operation *op, TensorType resultType, Value self,
+                               Value min, Value max,
+                               ConversionPatternRewriter &rewriter) {
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), self, min).failed())
+    return rewriter.notifyMatchFailure(op, "failed to equalize self and min");
+
+  auto selfType = cast<RankedTensorType>(self.getType());
+  auto minType = cast<RankedTensorType>(min.getType());
+
+  self = tosa::tosaCastTensorToType(rewriter, self,
+                                    selfType.clone(resultType.getElementType()))
+             .value();
+  min = tosa::tosaCastTensorToType(rewriter, min,
+                                   minType.clone(resultType.getElementType()))
+            .value();
+
+  auto maxRank = cast<RankedTensorType>(self.getType()).getRank();
+  auto dynamicIntermediateType =
+      RankedTensorType::get(SmallVector<int64_t>(maxRank, ShapedType::kDynamic),
+                            resultType.getElementType());
+
+  // max(xi, min_valuei)
+  // Use default NaN Propagation mode "PROPAGATE" for tosa.maximum
+  auto minThresholdCheck = tosa::CreateOpAndInfer<tosa::MaximumOp>(
+      rewriter, op->getLoc(), dynamicIntermediateType, self, min,
+      tosa::NanPropagationModeAttr::get(rewriter.getContext(),
+                                        tosa::NanPropagationMode::PROPAGATE));
+
+  Value tmp = minThresholdCheck.getResult();
+
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), tmp, max).failed())
+    return rewriter.notifyMatchFailure(
+        op, "failed to equalize intermediate and max");
+
+  auto maxType = cast<RankedTensorType>(max.getType());
+  max = tosa::tosaCastTensorToType(rewriter, max,
+                                   maxType.clone(resultType.getElementType()))
+            .value();
+
+  // yi = min(max(xi, min_valuei), max_valuei)
+  // Use default NaN Propagation mode "PROPAGATE" for tosa.minimum
+  auto result = tosa::CreateOpAndInfer<tosa::MinimumOp>(
+      rewriter, op->getLoc(), resultType, tmp, max,
+      tosa::NanPropagationModeAttr::get(rewriter.getContext(),
+                                        tosa::NanPropagationMode::PROPAGATE));
+
+  rewriter.replaceOp(op, result);
+  return success();
+}
+
+static LogicalResult validateClampBoundsInValidRange(
+    ConversionPatternRewriter &rewriter, Operation *op, IntegerType intType,
+    std::optional<int64_t> minInt, std::optional<int64_t> maxInt) {
+  auto isOutOfRange = [&](int64_t value) {
+    switch (intType.getWidth()) {
+    case 8:
+      return value < std::numeric_limits<int8_t>::min() ||
+             value > std::numeric_limits<int8_t>::max();
+    case 16:
+      return value < std::numeric_limits<int16_t>::min() ||
+             value > std::numeric_limits<int16_t>::max();
+    case 32:
+      return value < std::numeric_limits<int32_t>::min() ||
+             value > std::numeric_limits<int32_t>::max();
+    case 64:
+      return false;
+    default:
+      return true;
+    }
+  };
+
+  if ((minInt && isOutOfRange(*minInt)) || (maxInt && isOutOfRange(*maxInt))) {
+    return rewriter.notifyMatchFailure(
+        op, "explicit clamp bound is not representable in result integer type");
+  }
+
+  return success();
+}
+
 template <>
 LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewriteImpl(
     AtenClampOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
-
+  auto self = adaptor.getSelf();
   // Not a tensor type.
-  auto selfType = dyn_cast<TensorType>(adaptor.getSelf().getType());
+  auto selfType = dyn_cast<TensorType>(self.getType());
   if (!selfType)
     return rewriter.notifyMatchFailure(
         op, "only tensor types input are currently supported");
@@ -6013,7 +6093,6 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewriteImpl(
     return rewriter.notifyMatchFailure(
         op, "only tensor types output are currently supported");
   auto outElemTy = outType.getElementType();
-  Value self = adaptor.getSelf();
   if (selfType != outType) {
     auto castedSelf = tosa::tosaCastTensorToType(rewriter, self, outType);
     if (!castedSelf)
@@ -6055,34 +6134,76 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewriteImpl(
     }
   }
 
-  if (!isa<mlir::FloatType>(outElemTy)) {
-    IntegerAttr minIntAttr, maxIntAttr;
-    if (failed(tosa::getIntegerClampAttrs(rewriter, op, outElemTy, minInt,
-                                          maxInt, minIntAttr, maxIntAttr))) {
-      return failure();
-    }
-
-    rewriter.replaceOpWithNewOp<tosa::ClampOp>(
-        op, outType, self, minIntAttr, maxIntAttr,
-        /*nan_mode=*/
-        tosa::NanPropagationModeAttr::get(rewriter.getContext(),
-                                          tosa::NanPropagationMode::PROPAGATE));
-  } else {
-    FloatAttr minFloatAttr, maxFloatAttr;
-    if (failed(tosa::getFloatClampAttrs(rewriter, op, outElemTy, minFloat,
-                                        maxFloat, minFloatAttr,
-                                        maxFloatAttr))) {
-      return failure();
-    }
+  return TypeSwitch<Type, LogicalResult>(outElemTy)
+      .Case<mlir::IntegerType>([&](auto intType) -> LogicalResult {
+        if (failed(validateClampBoundsInValidRange(rewriter, op, intType,
+                                                   minInt, maxInt)))
+          return failure();
 
-    rewriter.replaceOpWithNewOp<tosa::ClampOp>(
-        op, outType, self, minFloatAttr, maxFloatAttr,
-        /*nan_mode=*/
-        tosa::NanPropagationModeAttr::get(rewriter.getContext(),
-                                          tosa::NanPropagationMode::PROPAGATE));
-  }
+        IntegerAttr minIntAttr, maxIntAttr;
+        if (failed(tosa::getIntegerClampAttrs(rewriter, op, outElemTy, minInt,
+                                              maxInt, minIntAttr,
+                                              maxIntAttr))) {
+          return failure();
+        }
+        // tosa.clamp does not support integer tensors wider than 16 bits.
+        //
+        // We use the following formula for 32-bit and 64-bit integers:
+        //    yi = min(max(xi, min_valuei), max_valuei)
+        switch (intType.getWidth()) {
+        case 8:
+        case 16:
+          if (minIntAttr.getInt() > maxIntAttr.getInt())
+            minIntAttr = maxIntAttr;
+          rewriter.replaceOpWithNewOp<tosa::ClampOp>(
+              op, outType, self, minIntAttr, maxIntAttr,
+              /*nan_mode=*/
+              tosa::NanPropagationModeAttr::get(
+                  rewriter.getContext(), tosa::NanPropagationMode::PROPAGATE));
+          return success();
+        case 32: {
+          int32_t minValue = static_cast<int32_t>(minIntAttr.getInt());
+          int32_t maxValue = static_cast<int32_t>(maxIntAttr.getInt());
+          Value min =
+              tosa::getConstTensor<int32_t>(rewriter, op, minValue, {}).value();
+          Value max =
+              tosa::getConstTensor<int32_t>(rewriter, op, maxValue, {}).value();
+          return rewriteClampAsMinimumMaximumOp(op, outType, self, min, max,
+                                                rewriter);
+        }
+        case 64: {
+          int64_t minValue = static_cast<int64_t>(minIntAttr.getInt());
+          int64_t maxValue = static_cast<int64_t>(maxIntAttr.getInt());
+          Value min =
+              tosa::getConstTensor<int64_t>(rewriter, op, minValue, {}).value();
+          Value max =
+              tosa::getConstTensor<int64_t>(rewriter, op, maxValue, {}).value();
+          return rewriteClampAsMinimumMaximumOp(op, outType, self, min, max,
+                                                rewriter);
+        }
+        default:
+          return rewriter.notifyMatchFailure(op, "Unsupported integer width");
+        }
+      })
+      .Case<mlir::FloatType>([&](auto) -> LogicalResult {
+        FloatAttr minFloatAttr, maxFloatAttr;
+        if (failed(tosa::getFloatClampAttrs(rewriter, op, outElemTy, minFloat,
+                                            maxFloat, minFloatAttr,
+                                            maxFloatAttr))) {
+          return failure();
+        }
 
-  return success();
+        rewriter.replaceOpWithNewOp<tosa::ClampOp>(
+            op, outType, self, minFloatAttr, maxFloatAttr,
+            /*nan_mode=*/
+            tosa::NanPropagationModeAttr::get(
+                rewriter.getContext(), tosa::NanPropagationMode::PROPAGATE));
+        return success();
+      })
+      .Default([&](Type) -> LogicalResult {
+        return rewriter.notifyMatchFailure(op,
+                                           "unsupported clamp element type");
+      });
 }
 
 // Legalization for aten.clamp.Tensor
@@ -6106,6 +6227,8 @@ LogicalResult ConvertAtenOp<AtenClampTensorOp>::matchAndRewriteImpl(
 
   auto resultType =
       dyn_cast<TensorType>(typeConverter->convertType(op.getType()));
+  if (!resultType)
+    return rewriter.notifyMatchFailure(op, "expected tensor result type");
 
   // Get min tensor. If None, there is no lower bound.
   Value min;
@@ -6126,6 +6249,11 @@ LogicalResult ConvertAtenOp<AtenClampTensorOp>::matchAndRewriteImpl(
                 return tosa::getConstTensor<int8_t>(
                            rewriter, op, std::numeric_limits<int8_t>::min(), {})
                     .value();
+              case 16:
+                return tosa::getConstTensor<int16_t>(
+                           rewriter, op, std::numeric_limits<int16_t>::min(),
+                           {})
+                    .value();
               case 32:
                 return tosa::getConstTensor<int32_t>(
                            rewriter, op, std::numeric_limits<int32_t>::min(),
@@ -6160,6 +6288,11 @@ LogicalResult ConvertAtenOp<AtenClampTensorOp>::matchAndRewriteImpl(
                 return tosa::getConstTensor<int8_t>(
                            rewriter, op, std::numeric_limits<int8_t>::max(), {})
                     .value();
+              case 16:
+                return tosa::getConstTensor<int16_t>(
+                           rewriter, op, std::numeric_limits<int16_t>::max(),
+                           {})
+                    .value();
               case 32:
                 return tosa::getConstTensor<int32_t>(
                            rewriter, op, std::numeric_limits<int32_t>::max(),
@@ -6175,33 +6308,8 @@ LogicalResult ConvertAtenOp<AtenClampTensorOp>::matchAndRewriteImpl(
             });
   }
 
-  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), self, min).failed() ||
-      mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), self, max).failed())
-    return rewriter.notifyMatchFailure(
-        op, "Failed to equalize ranks among operands and result");
-
-  self = tosa::tosaCastTensorToType(rewriter, self, resultType).value();
-  min = tosa::tosaCastTensorToType(rewriter, min, resultType).value();
-  max = tosa::tosaCastTensorToType(rewriter, max, resultType).value();
-
-  // max(xi, min_valuei)
-  // Use default NaN Propagation mode "PROPAGATE" for tosa.maximum
-  auto minThresholdCheck = tosa::MaximumOp::create(
-      rewriter, op->getLoc(), resultType, self, min,
-      /*nan_mode=*/
-      tosa::NanPropagationModeAttr::get(rewriter.getContext(),
-                                        tosa::NanPropagationMode::PROPAGATE));
-
-  // yi = min(max(xi, min_valuei), max_valuei)
-  // Use default NaN Propagation mode "PROPAGATE" for tosa.minimum
-  auto result = tosa::MinimumOp::create(
-      rewriter, op->getLoc(), resultType, minThresholdCheck, max,
-      /*nan_mode=*/
-      tosa::NanPropagationModeAttr::get(rewriter.getContext(),
-                                        tosa::NanPropagationMode::PROPAGATE));
-
-  rewriter.replaceOp(op, result);
-  return success();
+  return rewriteClampAsMinimumMaximumOp(op, resultType, self, min, max,
+                                        rewriter);
 }
 
 template <>

projects/pt1/e2e_testing/xfail_sets.py

@@ -26,6 +26,7 @@
     "ElementwiseAtanTensorBFloat16SpecialValuesModule_basic",
     "ElementwiseFloatTensorGtIntTensorModule_basic",
     "ElementwiseClampIntToFloatModule_basic",
+    "ElementwiseClampInt16Module_basic",
     # TODO: The values are extremely close to the golden values, but the test fails because of strict rtol/atol.
     "AtenInstanceNormModuleFp16_basic",
     "AtenIntMM_basic",
@@ -518,6 +519,7 @@
     "ReflectionPad3dModuleFront_basic",
     "ReflectionPad3dModuleBack_basic",
     "ElementwiseClampIntToFloatModule_basic",
+    "ElementwiseClampInt16Module_basic",
     # error: argument must be a memref of f32, f64, i32, i64, i8, i1, c32, c64, but got 'memref<3x5xbf16>'
     "ElementwiseClampMaxModule_bfloat16",
     "ElementwiseClampMinModule_bfloat16",
@@ -554,6 +556,7 @@
     "DiagonalWithStaticShapeModule_basic",
     "EinsumStaticDiagonalDimensionModule_basic",
     "ElementwiseAtanTensorBFloat16SpecialValuesModule_basic",
+    "ElementwiseClampInt16Module_basic",
     "ElementwiseClampIntToFloatModule_basic",
     "ElementwiseRemainderScalarModule_Bool_NegativeDivisor_basic",
     "ElementwiseRemainderScalarModule_Float_NegativeDividend_basic",
@@ -2249,9 +2252,12 @@
     "ElementwiseCeluModule_basic",
     "ElementwiseCeluStaticModule_basic",
     "ElementwiseClampMaxModule_basic",
+    "ElementwiseClampInt8MinGreaterThanMaxModule_basic",
     "ElementwiseClampMinModule_basic",
     "ElementwiseClampModule_basic",
     "ElementwiseClampTensorInt8Module_basic",
+    "ElementwiseClampInt32Module_basic",
+    "ElementwiseClampInt64Module_basic",
     "ElementwiseCloneChannelsLastMemoryFormatModule_basic",
     "ElementwiseCloneContiguousModule_basic",
     "ElementwiseCloneModule_basic",
@@ -3550,6 +3556,7 @@
     ONNX_XFAIL_SET = ONNX_XFAIL_SET | {
         "Aten_CastLongModule_basic",
         "Aten_CastFloatModule_basic",
+        "ElementwiseClampInt16Module_basic",
     }
 
 if torch_version_for_comparison() < version.parse("2.4.0.dev"):
@@ -4056,6 +4063,7 @@
     "ElementwiseClampMinModule_bfloat16",
     "ElementwiseClampModule_bfloat16",
     "ElementwiseReluModule_bfloat16",
+    "ElementwiseClampInt16Module_basic",  # 'memref<3x5xi16>'
 }
 
 ONNX_TOSA_CRASHING_SET = {
@@ -5195,4 +5203,5 @@
         # error: 'memref.cast' op operand type 'memref<2x6x4x3xf32>' and result type 'memref<2x6x5x3xf32>' are cast incompatible
         # torch.onnx.export produces onnx.MaxPool op with incorrect output shape of 2x6x5x3 instead of 2x6x4x3
         "MaxPool2dStaticCeilModeTrueReduceOutputModule_basic",
+        "ElementwiseClampInt16Module_basic",
     }