[Mlir-commits] [mlir] [mlir][linalg] Align `elementwise` builder to do type-conversion of input to result type (PR #190566)

[Julian Oppermann]

================
@@ -4947,14 +4950,56 @@ void ElementwiseOp::regionBuilder(
   SmallVector<Value> yields;
   Value result;
 
+  // Cast input value to dst type.
+  // Only same-kind casts are valid (float to float, int to int).
+  auto castToDstType = [&](Value v, Type dstType) -> Value {
+    Type srcType = v.getType();
+    if (srcType == dstType)
+      return v;
+
+    // Float -> Float
+    if (auto srcFloatType = dyn_cast<FloatType>(srcType)) {
+      if (auto dstFloatType = dyn_cast<FloatType>(dstType)) {
+        if (srcFloatType.getWidth() < dstFloatType.getWidth())
+          return arith::ExtFOp::create(b, b.getLoc(), dstType, v).getResult();
+        return arith::TruncFOp::create(b, b.getLoc(), dstType, v).getResult();
+      }
+    }
+
+    // Int -> Int
+    if (auto srcIntType = dyn_cast<IntegerType>(srcType)) {
+      if (auto dstIntType = dyn_cast<IntegerType>(dstType)) {
+        if (srcIntType.getWidth() < dstIntType.getWidth()) {
+          return srcIntType.isUnsigned()
+                     ? arith::ExtUIOp::create(b, b.getLoc(), dstType, v)
+                           .getResult()
+                     : arith::ExtSIOp::create(b, b.getLoc(), dstType, v)
+                           .getResult();
+        }
+        return arith::TruncIOp::create(b, b.getLoc(), dstType, v);
+      }
+    }
+
+    emitError() << "invalid cast from " << srcType << " to " << dstType
+                << " in linalg.elementwise";
+    return nullptr;
+  };
+
+  // Infer the compute element type from result type.
+  Type computeElementType = block.getArguments().back().getType();
+
+  // Create the linalg.generic body.
   if (arityGroup == ElementwiseArityGroup::Unary) {
-    result = helper.buildUnaryFn(kind.unaryFn, block.getArgument(0));
+    Value in0 = castToDstType(block.getArgument(0), computeElementType);
+    result = helper.buildUnaryFn(kind.unaryFn, in0);
 
   } else if (arityGroup == ElementwiseArityGroup::Binary) {
-    result = helper.buildBinaryFn(kind.binaryFn, block.getArgument(0),
-                                  block.getArgument(1));
+    Value in0 = castToDstType(block.getArgument(0), computeElementType);
+    Value in1 = castToDstType(block.getArgument(1), computeElementType);
+    result = helper.buildBinaryFn(kind.binaryFn, in0, in1);
 
   } else if (arityGroup == ElementwiseArityGroup::Ternary) {
+    // ternary op (select) should not be type casted.
----------------
jopperm wrote:

This exception makes sense to me, but again, should this be mentioned in the docs?

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