[MLIR][TORCH] Add E2E support for aten.masked_fill.Scalar op

This commit adds lowering of `aten.masked_fill.Scalar` op.
This commit also fixes the formatting of the file constant_alloc.py.

Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>

lib/Conversion/TorchToLinalg/Uncategorized.cpp

@@ -894,6 +894,18 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
                                           threshold);
     return b.create<arith::SelectOp>(loc, predicate, constantZero, grad);
   }
+  if (auto maskedFill = dyn_cast<AtenMaskedFillScalarOp>(op)) {
+    AtenMaskedFillScalarOp::Adaptor adaptor(operands);
+    Type dtype = converter->convertType(maskedFill.getType())
+                     .cast<RankedTensorType>()
+                     .getElementType();
+
+    Value input = payloadArgs[0];
+    Value mask = payloadArgs[1];
+    Value fillValue = convertScalarToDtype(b, loc, adaptor.value(), dtype);
+
+    return b.create<arith::SelectOp>(loc, mask, fillValue, input);
+  }
 
   op->emitError("unimplemented lowering in "
                 "createLinalgPayloadCalculationForElementwiseOp");
@@ -939,7 +951,7 @@ public:
              AtenWhereSelfOp, AtenCeilOp, AtenGtTensorOp, AtenEqTensorOp,
              AtenLtTensorOp, AtenSubScalarOp, AtenAddScalarOp, AtenThresholdOp,
              AtenThresholdBackwardOp, AtenCloneOp, AtenSinOp, AtenCosOp,
-             AtenNeScalarOp, AtenNegOp>(op))
+             AtenNeScalarOp, AtenNegOp, AtenMaskedFillScalarOp>(op))
       return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
 
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
@@ -1657,7 +1669,7 @@ void mlir::torch::torch_to_linalg::populateUncategorizedPatternsAndLegality(
       AtenGtScalarOp, AtenGeScalarOp, AtenEqScalarOp, AtenLtScalarOp,
       AtenLeScalarOp, AtenWhereSelfOp, AtenGtTensorOp, AtenEqTensorOp,
       AtenLtTensorOp, AtenThresholdOp, AtenThresholdBackwardOp, AtenCloneOp,
-      AtenSinOp, AtenCosOp, AtenNeScalarOp>();
+      AtenSinOp, AtenCosOp, AtenNeScalarOp, AtenMaskedFillScalarOp>();
   patterns.add<ConvertElementwiseOp>(typeConverter, context);
   target.addIllegalOp<AtenNllLossForwardOp>();
   patterns.add<ConvertAtenNllLossForwardOp>(typeConverter, context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -499,7 +499,7 @@ ChangeResult TypeAnalyzer::visitOperation(
           AtenRepeatOp, AtenConstantPadNdOp, AtenPadOp, AtenZero_Op,
           AtenIndexTensorOp, ValsemVariantAtenIndexPutImplOp, AtenIndexPutOp,
           ValsemVariantAtenCopyOp, ValsemVariantAtenZeroOp,
-          AtenIndexPutHackedTwinOp>(op)) {
+          AtenIndexPutHackedTwinOp, AtenMaskedFillScalarOp>(op)) {
     ValueKnowledge knowledge =
         ValueKnowledge::getNotNonePessimisticValueState(op->getContext());
     knowledge.dtype = operands[0]->getValue().dtype;

lib/Dialect/Torch/Transforms/ShapeLibrary.cpp

@@ -1981,6 +1981,10 @@ module {
     %0 = call @__torch__.torch_mlir.dialects.torch.importer.jit_ir.build_tools.upstream_shape_helpers.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>
     return %0 : !torch.list<int>
   }
+  func @"__torch_mlir_shape_fn.aten.masked_fill.Scalar"(%arg0: !torch.list<int>, %arg1: !torch.list<int>, %arg2: !torch.float) -> !torch.list<int> {
+    %0 = call @__torch__.torch_mlir.dialects.torch.importer.jit_ir.build_tools.upstream_shape_helpers.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>
+    return %0 : !torch.list<int>
+  }
   func @"__torch_mlir_shape_fn.aten.zero"(%arg0: !torch.list<int>) -> !torch.list<int> {
     return %arg0 : !torch.list<int>
   }
@@ -2047,7 +2051,7 @@ module {
       torch.prim.If.yield
     }
     %2 = torch.aten.sub.float %arg1, %arg0 : !torch.float, !torch.float -> !torch.float
-    %3 = torch.operator "aten.div.float"(%2, %arg2) : (!torch.float, !torch.float) -> !torch.float
+    %3 = torch.aten.div.float %2, %arg2 : !torch.float, !torch.float -> !torch.float
     %4 = torch.operator "aten.ceil.float"(%3) : (!torch.float) -> !torch.int
     %5 = torch.prim.ListConstruct %4 : (!torch.int) -> !torch.list<int>
     return %5 : !torch.list<int>

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py

@@ -625,6 +625,9 @@ def aten〇new_empty(self: List[int], size: List[int], dtype: Optional[int] = No
 def aten〇_to_copy(self: List[int], dtype: Optional[int] = None, layout: Optional[int] = None, device: Optional[device] = None, pin_memory: Optional[bool] = None, non_blocking: bool = False, memory_format: Optional[int] = None) -> List[int]:
     return upstream_shape_helpers.unary(self)
 
+def aten〇masked_fill〇Scalar(self: List[int], mask: List[int], value: float) -> List[int]:
+    return upstream_shape_helpers.unary(self)
+
 @not_present_in_registry
 def aten〇zero(self: List[int]) -> List[int]:
     return self