[TOSA] Add F16 and BF16 support for tosa.clamp

Signed-off-by: Justin Ngo <[email protected]>
Change-Id: I0f1ed0249018b40d068af3a2d02a93776c0c4a5f

Author: justin-ngo-arm

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -31,6 +31,20 @@
 
 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
 
+namespace {
+
+namespace limits {
+// Float 16 limits
+constexpr float Float16Max = 65504.0f;
+constexpr float Float16Lowest = -65504.0f;
+
+// BFloat 16 limits
+constexpr float BFloat16Max = 3.38953139e38f;
+constexpr float BFloat16Lowest = -3.38953139e38f;
+} // namespace limits
+
+} // namespace
+
 using namespace mlir;
 using namespace mlir::torch;
 using namespace mlir::torch::Torch;
@@ -871,6 +885,9 @@ LogicalResult ConvertAtenOp<AtenReluOp>::matchAndRewrite(
     ConversionPatternRewriter &rewriter) const {
   Value self = adaptor.getSelf();
   auto selfTy = cast<TensorType>(self.getType());
+  auto outType =
+      dyn_cast<TensorType>(getTypeConverter()->convertType(op.getType()));
+  auto outElemTy = outType.getElementType();
 
   if (!selfTy) {
     return rewriter.notifyMatchFailure(op,
@@ -883,12 +900,28 @@ LogicalResult ConvertAtenOp<AtenReluOp>::matchAndRewrite(
         op, "Only floating-point datatype legalization currently supported");
   }
 
+  FloatAttr minFloatAttr, maxFloatAttr;
+  if (outElemTy.isF16()) {
+    minFloatAttr = rewriter.getF16FloatAttr(0.0f);
+    maxFloatAttr = rewriter.getF16FloatAttr(limits::Float16Max);
+  } else if (outElemTy.isBF16()) {
+    minFloatAttr = rewriter.getFloatAttr(rewriter.getBF16Type(), 0.0f);
+    maxFloatAttr =
+        rewriter.getFloatAttr(rewriter.getBF16Type(), limits::BFloat16Max);
+  } else if (outElemTy.isF32()) {
+    minFloatAttr = rewriter.getF32FloatAttr(0.0f);
+    maxFloatAttr = rewriter.getF32FloatAttr(std::numeric_limits<float>::max());
+  } else if (outElemTy.isF64()) {
+    minFloatAttr = rewriter.getF64FloatAttr(0.0f);
+    maxFloatAttr = rewriter.getF64FloatAttr(std::numeric_limits<double>::max());
+  } else {
+    return rewriter.notifyMatchFailure(op, "Unsupported floating-point type");
+  }
+
   // Maps to tosa.clamp
   // Use default NaN Propagation mode "PROPAGATE" for tosa.clamp
   rewriter.replaceOpWithNewOp<tosa::ClampOp>(
-      op, getTypeConverter()->convertType(op.getType()), self,
-      rewriter.getF32FloatAttr(0.0f),
-      rewriter.getF32FloatAttr(std::numeric_limits<float>::max()),
+      op, outType, self, minFloatAttr, maxFloatAttr,
       /*nan_mode=*/rewriter.getStringAttr("PROPAGATE"));
   return success();
 }
@@ -5186,10 +5219,32 @@ LogicalResult ConvertAtenOp<AtenClampOp>::matchAndRewrite(
         op, outType, adaptor.getSelf(), minIntAttr, maxIntAttr,
         /*nan_mode=*/rewriter.getStringAttr("PROPAGATE"));
   } else {
-    FloatAttr minFloatAttr = rewriter.getF32FloatAttr(
-        isMinNotNone ? minFloat : std::numeric_limits<float>::lowest());
-    FloatAttr maxFloatAttr = rewriter.getF32FloatAttr(
-        isMaxNotNone ? maxFloat : std::numeric_limits<float>::max());
+    FloatAttr minFloatAttr, maxFloatAttr;
+    if (outElemTy.isF16()) {
+      minFloatAttr = rewriter.getF16FloatAttr(
+          isMinNotNone ? minFloat : limits::Float16Lowest);
+      maxFloatAttr = rewriter.getF16FloatAttr(
+          isMaxNotNone ? maxFloat : limits::Float16Max);
+    } else if (outElemTy.isBF16()) {
+      minFloatAttr = rewriter.getFloatAttr(
+          rewriter.getBF16Type(),
+          isMinNotNone ? minFloat : limits::BFloat16Lowest);
+      maxFloatAttr =
+          rewriter.getFloatAttr(rewriter.getBF16Type(),
+                                isMaxNotNone ? maxFloat : limits::BFloat16Max);
+    } else if (outElemTy.isF32()) {
+      minFloatAttr = rewriter.getF32FloatAttr(
+          isMinNotNone ? minFloat : std::numeric_limits<float>::lowest());
+      maxFloatAttr = rewriter.getF32FloatAttr(
+          isMaxNotNone ? maxFloat : std::numeric_limits<float>::max());
+    } else if (outElemTy.isF64()) {
+      minFloatAttr = rewriter.getF64FloatAttr(
+          isMinNotNone ? minFloat : std::numeric_limits<double>::lowest());
+      maxFloatAttr = rewriter.getF64FloatAttr(
+          isMaxNotNone ? maxFloat : std::numeric_limits<double>::max());
+    } else {
+      return rewriter.notifyMatchFailure(op, "Unsupported floating-point type");
+    }
 
     rewriter.replaceOpWithNewOp<tosa::ClampOp>(
         op, outType, adaptor.getSelf(), minFloatAttr, maxFloatAttr,
@@ -8547,14 +8602,29 @@ LogicalResult ConvertAtenOp<AtenLogitOp>::matchAndRewrite(
 
   auto zi = self;
 
+  FloatAttr minFloatAttr, maxFloatAttr;
+  if (resultElemTy.isF16()) {
+    minFloatAttr = rewriter.getF16FloatAttr(eps);
+    maxFloatAttr = rewriter.getF16FloatAttr(1 - eps);
+  } else if (resultElemTy.isBF16()) {
+    minFloatAttr = rewriter.getFloatAttr(rewriter.getBF16Type(), eps);
+    maxFloatAttr = rewriter.getFloatAttr(rewriter.getBF16Type(), 1 - eps);
+  } else if (resultElemTy.isF32()) {
+    minFloatAttr = rewriter.getF32FloatAttr(eps);
+    maxFloatAttr = rewriter.getF32FloatAttr(1 - eps);
+  } else if (resultElemTy.isF64()) {
+    minFloatAttr = rewriter.getF64FloatAttr(eps);
+    maxFloatAttr = rewriter.getF64FloatAttr(1 - eps);
+  } else {
+    return rewriter.notifyMatchFailure(op, "Unsupported floating-point type");
+  }
+
   // Clamp input to [eps, 1 - eps] when eps is not None
   // Use default NaN Propagation mode "PROPAGATE" for tosa.clamp
   if (!isEpsNone) {
     zi = rewriter
              .create<tosa::ClampOp>(
-                 op->getLoc(), resultType, self,
-                 rewriter.getF32FloatAttr(static_cast<float>(eps)),
-                 rewriter.getF32FloatAttr(static_cast<float>(1 - eps)),
+                 op->getLoc(), resultType, self, minFloatAttr, maxFloatAttr,
                  /*nan_mode=*/rewriter.getStringAttr("PROPAGATE"))
              .getResult();
   }