[Mlir-commits] [mlir] [mlir][tosa] Avoid introducing int <-> float casts (PR #195882)

[Ian Tayler Lessa]

================
@@ -907,40 +907,98 @@ struct NonNarrowingCastsOptimization : public OpRewritePattern<tosa::CastOp> {
 
     const Value innerCastInput = innerCastOp.getInput();
 
-    const auto innerInputType =
+    const ShapedType innerInputType =
         llvm::cast<ShapedType>(innerCastInput.getType());
-    const auto innerOutputType = llvm::cast<ShapedType>(innerCastOp.getType());
-    const auto outerOutputType = llvm::cast<ShapedType>(castOp.getType());
+    const ShapedType innerOutputType =
+        llvm::cast<ShapedType>(innerCastOp.getType());
+    const ShapedType outerOutputType = llvm::cast<ShapedType>(castOp.getType());
 
-    const SmallVector<ShapedType, 3> types = {innerInputType, innerOutputType,
-                                              outerOutputType};
+    const Type innerInputElemType = innerInputType.getElementType();
+    const Type innerOutputElemType = innerOutputType.getElementType();
+    const Type outerOutputElemType = outerOutputType.getElementType();
 
-    if (llvm::any_of(types, [](const ShapedType type) {
-          const auto elemTy = type.getElementType();
+    const SmallVector<Type, 3> types = {innerInputElemType, innerOutputElemType,
+                                        outerOutputElemType};
+
+    if (llvm::any_of(types, [](const Type type) {
           // Support a specific set of floating point types since we need to be
           // careful in not introducing unsupported type combinations
-          return !(elemTy.isInteger() ||
+          return !(type.isInteger() ||
                    llvm::isa<Float8E4M3FNType, Float8E5M2Type, BFloat16Type,
-                             Float16Type, Float32Type>(elemTy));
+                             Float16Type, Float32Type>(type));
         }))
       return rewriter.notifyMatchFailure(
           castOp, "only integer and f32, f16, bf16, f8E4M3FN, f8E5M2 types are "
                   "supported");
 
-    if (llvm::isa<Float8E5M2Type>(innerInputType.getElementType()) &&
-        llvm::isa<Float8E4M3FNType>(outerOutputType.getElementType())) {
+    if (llvm::isa<Float8E5M2Type>(innerInputElemType) &&
+        llvm::isa<Float8E4M3FNType>(outerOutputElemType)) {
       return rewriter.notifyMatchFailure(
           castOp, "avoid introducing f8E5M2 -> f8E4M3FN casts which are not "
                   "legal in TOSA");
     }
 
-    if (llvm::isa<Float8E4M3FNType>(innerInputType.getElementType()) &&
-        llvm::isa<Float8E5M2Type>(outerOutputType.getElementType())) {
+    if (llvm::isa<Float8E4M3FNType>(innerInputElemType) &&
+        llvm::isa<Float8E5M2Type>(outerOutputElemType)) {
       return rewriter.notifyMatchFailure(
           castOp, "avoid introducing f8E4M3FN -> f8E5M2 casts which are not "
                   "legal in TOSA");
     }
 
+    if (llvm::isa<Float8E5M2Type, Float8E4M3FNType>(innerInputElemType) &&
+        outerOutputElemType.isInteger()) {
+      return rewriter.notifyMatchFailure(
+          castOp, "avoid introducing fp8 -> integer casts which are not "
+                  "legal in TOSA");
+    }
+
+    if (innerInputElemType.isInteger() &&
+        llvm::isa<Float8E5M2Type, Float8E4M3FNType>(outerOutputElemType)) {
+      return rewriter.notifyMatchFailure(
+          castOp, "avoid introducing integer -> fp8 casts which are not "
+                  "legal in TOSA");
+    }
+
+    if (llvm::isa<Float16Type>(innerInputElemType) &&
+        llvm::isa<BFloat16Type>(outerOutputElemType)) {
+      return rewriter.notifyMatchFailure(
+          castOp, "avoid introducing fp16 -> bf16 casts which are not "
+                  "legal in TOSA");
+    }
+
+    if (llvm::isa<BFloat16Type>(innerInputElemType) &&
+        llvm::isa<Float16Type>(outerOutputElemType)) {
+      return rewriter.notifyMatchFailure(
+          castOp, "avoid introducing bf16 -> fp16 casts which are not "
+                  "legal in TOSA");
+    }
+
+    const auto isIntegerOneOf = [](Type type, size_t bitwidth1,
----------------
IanTaylerLessa-arm wrote:

I considered the type table but I found that it would be a bit tricky to implement nicely. I think the nicer way to implement this is to expose some utilities from the validation headers (`TosaProfileCompliance`) that can then be used in passes.

The part making it tricky is that float types are checked by type, while integer types need to be checked against the bit-width, so one is a compile-time value and the other is runtime. We'd probably want to add an enum with all the types and a couple of utilities with switches that do 'the right thing' for each type in the enum, but then we're replicating a lot of the code on `TosaProfileCompliance` with `TypeInfo`, etc.

The bigger re-shuffling to expose that seemed out of scope for this PR, and replicating the logic also seemed less-than-ideal, so I went for the option of these more ad-hoc checks for now, with the idea that we could refactor it if we do expose these utilities later.

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