[Mlir-commits] [mlir] [mlir][tosa] Add clamp op support to `TosaNarrowI64ToI32` pass (PR #169308)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Luke Hutton (lhutton1)

<details>
<summary>Changes</summary>

This commit allows the narrowing of `tosa.clamp` when the min/max attributes are within the int32 range.

---
Full diff: https://github.com/llvm/llvm-project/pull/169308.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaNarrowI64ToI32.cpp (+38) 
- (modified) mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir (+9) 
- (modified) mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir (+9) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowI64ToI32.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowI64ToI32.cpp
index ddaf7d8a5e033..be442cc4f88ca 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowI64ToI32.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowI64ToI32.cpp
@@ -176,6 +176,42 @@ class ConvertCastOpWithBoundsChecking
   }
 };
 
+class ConvertClampOpWithBoundsChecking
+    : public OpConversionPattern<tosa::ClampOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(tosa::ClampOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const final {
+    const auto minAttr = dyn_cast<IntegerAttr>(op.getMinValAttr());
+    const auto maxAttr = dyn_cast<IntegerAttr>(op.getMaxValAttr());
+    if (!minAttr || !maxAttr)
+      return failure();
+
+    const int64_t min = minAttr.getInt();
+    const int64_t max = maxAttr.getInt();
+
+    if (min < std::numeric_limits<int32_t>::min() ||
+        max > std::numeric_limits<int32_t>::max())
+      return rewriter.notifyMatchFailure(
+          op, "Clamp bounds exceed int32 range. Narrowing cast may lead to "
+              "data loss.");
+
+    const Type resultType = op.getOutput().getType();
+    const Type newResultType = typeConverter->convertType(resultType);
+
+    const IntegerType int32Type = IntegerType::get(rewriter.getContext(), 32);
+    const IntegerAttr newMinAttr =
+        rewriter.getIntegerAttr(int32Type, static_cast<int32_t>(min));
+    const IntegerAttr newMaxAttr =
+        rewriter.getIntegerAttr(int32Type, static_cast<int32_t>(max));
+    rewriter.replaceOpWithNewOp<tosa::ClampOp>(op, newResultType,
+                                               adaptor.getInput(), newMinAttr,
+                                               newMaxAttr, op.getNanModeAttr());
+    return success();
+  }
+};
+
 template <typename OpTy>
 class ConvertTypedOp : public OpConversionPattern<OpTy> {
   using OpConversionPattern<OpTy>::OpConversionPattern;
@@ -285,6 +321,8 @@ struct TosaNarrowI64ToI32
     } else {
       // Tensor
       patterns.add<ConvertArgMaxOpWithBoundsChecking>(typeConverter, context);
+      // Activation functions
+      patterns.add<ConvertClampOpWithBoundsChecking>(typeConverter, context);
       // Data layout
       patterns.add<ConvertTypedOp<tosa::ConcatOp>>(typeConverter, context);
       patterns.add<ConvertTypedOp<tosa::PadOp>>(typeConverter, context);
diff --git a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir
index 1a36177a37033..9848fe4abb345 100644
--- a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir
@@ -79,3 +79,12 @@ func.func @test_const() -> tensor<2xi64> {
   // FUNCBOUND: return %[[CONST]] : tensor<2xi32>
   return %0 : tensor<2xi64>
 }
+
+// -----
+
+// CHECK-LABEL: test_clamp_trunc
+func.func @test_clamp_trunc(%arg0: tensor<100xi64>) -> tensor<100xi64> {
+  // COMMON: tosa.clamp %{{.*}} {max_val = 2147483647 : i32, min_val = -2147483648 : i32} : (tensor<100xi32>) -> tensor<100xi32>
+  %1 = tosa.clamp %arg0 {max_val = 3000000000 : i64, min_val = -2147483648 : i64} : (tensor<100xi64>) -> tensor<100xi64>
+  return %1 : tensor<100xi64>
+}
diff --git a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
index a14483fcdd7b0..c10393296a5df 100644
--- a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
@@ -160,3 +160,12 @@ func.func @test_transition_from_i64(%arg0: tensor<1xi64>) -> tensor<1xi32> {
   // COMMON: return %[[OUT_CAST]] : tensor<1xi32>
   return %2 : tensor<1xi32>
 }
+
+// -----
+
+// CHECK-LABEL: test_clamp
+func.func @test_clamp(%arg0: tensor<100xi64>) -> tensor<100xi64> {
+  // COMMON: tosa.clamp %{{.*}} {max_val = 2147483647 : i32, min_val = -2147483648 : i32} : (tensor<100xi32>) -> tensor<100xi32>
+  %1 = tosa.clamp %arg0 {max_val = 2147483647 : i64, min_val = -2147483648 : i64} : (tensor<100xi64>) -> tensor<100xi64>
+  return %1 : tensor<100xi64>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/169308