Add float support for masked_fill

Author: gpetters94

lib/Conversion/TorchToLinalg/Uncategorized.cpp

@@ -906,6 +906,8 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
 
     Value input = payloadArgs[0];
     Value mask = payloadArgs[1];
+    if (mask.getType().isa<mlir::FloatType>())
+      mask = b.create<arith::ConstantOp>(loc, b.getBoolAttr(false));
     Value fillValue = convertScalarToDtype(b, loc, adaptor.value(), dtype);
 
     return b.create<arith::SelectOp>(loc, mask, fillValue, input);

python/torch_mlir_e2e_test/test_suite/constant_alloc.py

@@ -1462,3 +1462,25 @@ def forward(self, x, mask, value):
 def MaskedFillTensorFloatValueModule_basic(module, tu: TestUtils):
     module.forward(tu.randint(2, 3, low=-10, high=10),
                    tu.randint(2, 3, high=2).to(dtype=torch.bool), tu.rand())
+
+
+class MaskedFillTensorFloatMaskModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.int64, True),
+        ([-1, -1], torch.float, True),
+    ])
+    def forward(self, x, mask):
+        return torch.ops.aten.masked_fill(x, mask, value=0.1)
+
+
+@register_test_case(module_factory=lambda: MaskedFillTensorFloatMaskModule())
+def MaskedFillTensorFloatMaskModule_basic(module, tu: TestUtils):
+    mask = tu.randint(2, 3, low=0, high=2).to(torch.float)
+    module.forward(tu.randint(2, 3, low=-10, high=10),
+                   mask)