The fillOp's value needs to casted

[pashu123]

Summary

During elementwise fusion the fillOp's value was directly
referred without casting which can create mismatching dtypes.

[pashu123]

@mravishankar I have copied the cast funtion from https://github.com/llvm/llvm-project/blob/254bf678d49185026d1b17c6345106f9f3fe58b2/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L423 do let me know if there is a better way.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

1730

Do you mind moving this into include/mlir/Dialect/Arith/Utils/Utils.h and implementation into lib/Dialect/Arith/Utils/Utils.cpp?

[pashu123]

Move the casting function to arith/utils.

[pashu123]

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

1730

Done, PTAL.

[pashu123]

fix typo.

[mravishankar]

Thanks!

[ThomasRaoux]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

83

Can you remove the cast function and use this helper instead:
https://github.com/llvm/llvm-project/blob/254bf678d49185026d1b17c6345106f9f3fe58b2/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L423

[ThomasRaoux]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

90

for linalg ops the type would be signless, therefore this is probably not the right way to decide between signed and unsigned cast

[pashu123]

Remove the cast duplicate function.

[ThomasRaoux]

Thanks!

mlir/include/mlir/Dialect/Arith/Utils/Utils.h

Value getValueOrCreateCastToIndexLike(OpBuilder &b, Location loc,
                                       Type targetType, Value value);
 
 /// Similar to the other overload, but converts multiple OpFoldResults into
 /// Values.
 SmallVector<Value>
 getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
                                 ArrayRef<OpFoldResult> valueOrAttrVec);
 
+/// Converts a scalar value `operand` to type `toType`. If the value doesn't
+/// convert, a warning will be issued and the operand is returned as is (which
+/// will presumably yield a verification issue downstream).
+Value convertScalarToDtype(OpBuilder &b, Location loc, Value operand,
+                           Type toType, bool isUnsignedCast);
+
 /// Helper struct to build simple arithmetic quantities with minimal type
 /// inference support.
 struct ArithBuilder {
   ArithBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
 
   Value _and(Value lhs, Value rhs);
   Value add(Value lhs, Value rhs);
   Value sub(Value lhs, Value rhs);

mlir/lib/Dialect/Arith/Utils/Utils.cpp

assert(targetIntegerType && valueIntegerType &&
          "unexpected cast between types other than integers and index");
   assert(targetIntegerType.getSignedness() == valueIntegerType.getSignedness());
 
   if (targetIntegerType.getWidth() > valueIntegerType.getWidth())
     return b.create<arith::ExtSIOp>(loc, targetIntegerType, value);
   return b.create<arith::TruncIOp>(loc, targetIntegerType, value);
 }
 
+Value mlir::convertScalarToDtype(OpBuilder &b, Location loc, Value operand,
+                                 Type toType, bool isUnsignedCast) {
+  if (operand.getType() == toType)
+    return operand;
+  if (auto toIntType = toType.dyn_cast<IntegerType>()) {
+    // If operand is floating point, cast directly to the int type.
+    if (operand.getType().isa<FloatType>()) {
+      if (isUnsignedCast)
+        return b.create<arith::FPToUIOp>(loc, toType, operand);
+      return b.create<arith::FPToSIOp>(loc, toType, operand);
+    }
+    // Cast index operands directly to the int type.
+    if (operand.getType().isIndex())
+      return b.create<arith::IndexCastOp>(loc, toType, operand);
+    if (auto fromIntType = operand.getType().dyn_cast<IntegerType>()) {
+      // Either extend or truncate.
+      if (toIntType.getWidth() > fromIntType.getWidth()) {
+        if (isUnsignedCast)
+          return b.create<arith::ExtUIOp>(loc, toType, operand);
+        return b.create<arith::ExtSIOp>(loc, toType, operand);
+      }
+      if (toIntType.getWidth() < fromIntType.getWidth())
+        return b.create<arith::TruncIOp>(loc, toType, operand);
+    }
+  } else if (auto toFloatType = toType.dyn_cast<FloatType>()) {
+    // If operand is integer, cast directly to the float type.
+    // Note that it is unclear how to cast from BF16<->FP16.
+    if (operand.getType().isa<IntegerType>()) {
+      if (isUnsignedCast)
+        return b.create<arith::UIToFPOp>(loc, toFloatType, operand);
+      return b.create<arith::SIToFPOp>(loc, toFloatType, operand);
+    }
+    if (auto fromFloatType = operand.getType().dyn_cast<FloatType>()) {
+      if (toFloatType.getWidth() > fromFloatType.getWidth())
+        return b.create<arith::ExtFOp>(loc, toFloatType, operand);
+      if (toFloatType.getWidth() < fromFloatType.getWidth())
+        return b.create<arith::TruncFOp>(loc, toFloatType, operand);
+    }
+  }
+  emitWarning(loc) << "could not cast operand of type " << operand.getType()
+                   << " to " << toType;
+  return operand;
+}
+
 SmallVector<Value>
 mlir::getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
                                       ArrayRef<OpFoldResult> valueOrAttrVec) {
   return llvm::to_vector<4>(
       llvm::map_range(valueOrAttrVec, [&](OpFoldResult value) -> Value {
         return getValueOrCreateConstantIndexOp(b, loc, value);
       }));
 }

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 private:
   // Generates operations to cast the given operand to a specified type.
   // If the cast cannot be performed, a warning will be issued and the
   // operand returned as-is (which will presumably yield a verification
   // issue downstream).
   Value cast(Type toType, Value operand, bool isUnsignedCast) {
     OpBuilder builder = getBuilder();
     auto loc = operand.getLoc();
-
-    if (operand.getType() == toType)
-      return operand;
-    if (auto toIntType = toType.dyn_cast<IntegerType>()) {
-      // If operand is floating point, cast directly to the int type.
-      if (operand.getType().isa<FloatType>()) {
-        if (isUnsignedCast)
-          return builder.create<arith::FPToUIOp>(loc, toType, operand);
-        return builder.create<arith::FPToSIOp>(loc, toType, operand);
-      }
-      // Cast index operands directly to the int type.
-      if (operand.getType().isIndex())
-        return builder.create<arith::IndexCastOp>(loc, toType, operand);
-      if (auto fromIntType = operand.getType().dyn_cast<IntegerType>()) {
-        // Either extend or truncate.
-        if (toIntType.getWidth() > fromIntType.getWidth()) {
-          if (isUnsignedCast)
-            return builder.create<arith::ExtUIOp>(loc, toType, operand);
-          return builder.create<arith::ExtSIOp>(loc, toType, operand);
-        }
-        if (toIntType.getWidth() < fromIntType.getWidth())
-          return builder.create<arith::TruncIOp>(loc, toType, operand);
-      }
-    } else if (auto toFloatType = toType.dyn_cast<FloatType>()) {
-      // If operand is integer, cast directly to the float type.
-      // Note that it is unclear how to cast from BF16<->FP16.
-      if (operand.getType().isa<IntegerType>()) {
-        if (isUnsignedCast)
-          return builder.create<arith::UIToFPOp>(loc, toFloatType, operand);
-        return builder.create<arith::SIToFPOp>(loc, toFloatType, operand);
-      }
-      if (auto fromFloatType = operand.getType().dyn_cast<FloatType>()) {
-        if (toFloatType.getWidth() > fromFloatType.getWidth())
-          return builder.create<arith::ExtFOp>(loc, toFloatType, operand);
-        if (toFloatType.getWidth() < fromFloatType.getWidth())
-          return builder.create<arith::TruncFOp>(loc, toFloatType, operand);
-      }
-    }
-
-    emitWarning(operand.getLoc()) << "could not cast operand of type "
-                                  << operand.getType() << " to " << toType;
-    return operand;
+    return convertScalarToDtype(builder, loc, operand, toType, isUnsignedCast);
   }
 
   bool isComplex(Value value) { return value.getType().isa<ComplexType>(); }
   bool isFloatingPoint(Value value) { return value.getType().isa<FloatType>(); }
   bool isInteger(Value value) { return value.getType().isa<IntegerType>(); }
 
   OpBuilder getBuilder() {
     OpBuilder builder(context);

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

if (!modifiedOutput) {
       rewriter.cancelRootUpdate(op);
       return failure();
     }
     rewriter.finalizeRootUpdate(op);
     return success();
   }
 };
 
 /// Fold linalg.fill into linalg.generic
 struct FoldFillWithGenericOp : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     if (!genericOp.hasTensorSemantics())
       return failure();
     bool fillFound = false;
     Block &payload = genericOp.getRegion().front();
     for (OpOperand *opOperand : genericOp.getDpsInputOperands()) {
       if (!genericOp.payloadUsesValueFromOperand(opOperand))
         continue;
       FillOp fillOp = opOperand->get().getDefiningOp<FillOp>();
       if (!fillOp)
         continue;
       fillFound = true;
+      Value fillVal = fillOp.value();
+      auto resultType =
+          fillOp.result().getType().cast<RankedTensorType>().getElementType();
+      Value convertedVal =
+          convertScalarToDtype(rewriter, fillOp.getLoc(), fillVal, resultType,
+                               /*isUnsignedCast =*/false);
       payload.getArgument(opOperand->getOperandNumber())
-          .replaceAllUsesWith(fillOp.value());
+          .replaceAllUsesWith(convertedVal);
     }
     return success(fillFound);
   }
 };
 } // namespace
 
 void mlir::linalg::populateFoldReshapeOpsByExpansionPatterns(
     RewritePatternSet &patterns,

mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir

^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
     %5 = arith.addf  %arg1, %arg2 : f32
         linalg.yield %5 : f32
   } -> tensor<?xf32>
   return %4 : tensor<?xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @fold_fill_generic_different_dtype
+//  CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf16>) -> tensor<?xf16> { 
+//   CHECK-NOT: linalg.fill
+//       CHECK: %[[GENERIC_OP:.*]] = linalg.generic
+//  CHECK-SAME: ins(%[[ARG0]] : tensor<?xf16>)
+//  CHECK-SAME: outs({{.*}} : tensor<?xf16>) {
+#map0 = affine_map<(d0) -> (d0)>
+func.func @fold_fill_generic_different_dtype(%arg0: tensor<?xf16>) -> (tensor<?xf16>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 7.0 : f32
+  %0 = tensor.dim %arg0, %c0 : tensor<?xf16>
+  %1 = tensor.empty(%0) : tensor<?xf16>
+  %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?xf16>) -> tensor<?xf16>
+  %3 = tensor.empty(%0) : tensor<?xf16>
+  %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %2 : tensor<?xf16>, tensor<?xf16>) outs (%3:tensor<?xf16>) {
+  ^bb0(%arg1: f16, %arg2: f16, %arg3: f16):
+    %5 = arith.addf  %arg1, %arg2 : f16
+        linalg.yield %5 : f16
+  } -> tensor<?xf16>
+  return %4 : tensor<?xf16>
+}
+
+// -----
+
 // CHECK-LABEL: func @fold_fill_generic_mixedaccess
 //   CHECK-NOT: linalg.fill
 //       CHECK: %[[GENERIC_OP:.*]] = linalg.generic
 //   CHECK-NOT: ins
 //  CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) {
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1) -> (d1, d0)>
 func.func @fold_fill_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) {