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 ↗

(On Diff #473421)

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 ↗

(On Diff #473421)

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/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

if (!modifiedOutput) {
       rewriter.cancelRootUpdate(op);
       return failure();
     }
     rewriter.finalizeRootUpdate(op);
     return success();
   }
 };
 
+// Casts the Value `operand` to type `toType`.
+static Value cast(PatternRewriter &builder, Type toType, Value operand) {
+  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 (toIntType.isUnsigned())
+        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 (toIntType.isUnsigned())
+          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 (auto intType = operand.getType().dyn_cast<IntegerType>()) {
+      if (intType.isUnsigned())
+        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 nullptr;
+}
+
 /// 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 castedVal = cast(rewriter, resultType, fillVal);
+      if (!castedVal)
+        return failure();
       payload.getArgument(opOperand->getOperandNumber())
-          .replaceAllUsesWith(fillOp.value());
+          .replaceAllUsesWith(castedVal);
     }
     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>) {