Add cases for view, type_as, and masked_fill

Author: gpetters94

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -6133,6 +6133,7 @@ def Torch_AtenTypeAsOp : Torch_Op<"aten.type_as", [
       printDefaultTorchOp(printer, *this, 2, 1);
     }
   }];
+  let hasFolder = 1;
 }
 
 def Torch_AtenViewOp : Torch_Op<"aten.view", [
@@ -6157,6 +6158,7 @@ def Torch_AtenViewOp : Torch_Op<"aten.view", [
     }
   }];
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Torch_Aten_UnsafeViewOp : Torch_Op<"aten._unsafe_view", [

lib/Conversion/TorchToLinalg/DataMovement.cpp

@@ -301,6 +301,42 @@ class ConvertAtenViewOp : public OpConversionPattern<AtenViewOp> {
           op, "desired size list length mismatches with the result type rank");
     }
 
+    // Case where all sizes (input and output) are statically known
+    SmallVector<int64_t> outputSizeIntList;
+    if (matchPattern(op.size(), m_TorchConstantIntList(outputSizeIntList)) &&
+        llvm::none_of(outputSizeIntList,
+                      [](int64_t x) { return x == kUnknownSize; }) &&
+        inputType.hasStaticShape()) {
+
+      int64_t midSize = 1;
+      for (auto outSize : outputSizeIntList)
+        midSize *= outSize;
+      auto midType =
+          RankedTensorType::get({midSize}, resultType.getElementType());
+      auto outType =
+          RankedTensorType::get(outputSizeIntList, resultType.getElementType());
+      int64_t outputRank = outputSizeIntList.size();
+
+      SmallVector<ReassociationIndices> flatten(1);
+      for (auto i = 0; i < inputRank; i++)
+        flatten[0].push_back(i);
+      SmallVector<ReassociationIndices> unflatten(1);
+      for (auto i = 0; i < outputRank; i++)
+        unflatten[0].push_back(i);
+
+      Value flattened =
+          (inputRank <= 1 ? input
+                          : rewriter.create<tensor::CollapseShapeOp>(
+                                loc, midType, input, flatten));
+      Value result =
+          (outputRank <= 1 ? flattened
+                           : rewriter.create<tensor::ExpandShapeOp>(
+                                 loc, outType, flattened, unflatten));
+
+      rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, result);
+      return success();
+    }
+
     // Currently, we only handle the cases where each dimension is either
     // being expanded or collapsed. We do not handle cases where it's neither
     // collapsing nor expanding like view of [2,3] for 3x2 tensor.
@@ -319,6 +355,7 @@ class ConvertAtenViewOp : public OpConversionPattern<AtenViewOp> {
     // collapsed. Note this may technically not always be true.
     // TODO: think of a way better way to at least detect when this assumption
     // is violated for the cases of dynamic dimensions.
+
     SmallVector<int64_t> outputShape(resultRank, kUnknownSize);
     SmallVector<ReassociationIndices> unchangedDims;
     llvm::Optional<int64_t> inferredDimension;

lib/Conversion/TorchToLinalg/Uncategorized.cpp

@@ -892,6 +892,9 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
 
     Value input = payloadArgs[0];
     Value mask = payloadArgs[1];
+    if (mask.getType().isa<mlir::FloatType>())
+      mask = b.create<arith::FPToUIOp>(
+          loc, IntegerType::get(op->getContext(), 1), mask);
     Value fillValue = convertScalarToDtype(b, loc, adaptor.value(), dtype);
 
     return b.create<arith::SelectOp>(loc, mask, fillValue, input);

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -673,6 +673,22 @@ OpFoldResult AtenSqueezeDimOp::fold(ArrayRef<Attribute> operands) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// AtenTypeAsOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AtenTypeAsOp::fold(ArrayRef<Attribute> operands) {
+  Value input = getOperand(0);
+  Value output = getOperand(1);
+  Type inType = input.getType();
+  Type outType = output.getType();
+
+  if (inType == outType)
+    return input;
+
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // AtenToDtypeOp
 //===----------------------------------------------------------------------===//
@@ -780,6 +796,36 @@ OpFoldResult AtenViewOp::fold(ArrayRef<Attribute> operands) {
   return getOperand(0);
 }
 
+void AtenViewOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                             MLIRContext *context) {
+  patterns.add(+[](AtenViewOp op, PatternRewriter &rewriter) {
+    Location loc = op->getLoc();
+    auto outType = op.getType().dyn_cast<BaseTensorType>();
+    Value size = op.size();
+    if (!outType || !outType.hasSizes())
+      return failure();
+    // When all result sizes are statically known, replace the sizes argument
+    // with a constant list
+    auto staticSizes = outType.getSizes();
+    SmallVector<int64_t> constSizeList;
+    if (!matchPattern(size, m_TorchConstantIntList(constSizeList)) &&
+        llvm::none_of(staticSizes,
+                      [](int64_t x) { return x == kUnknownSize; })) {
+      SmallVector<Value> staticSizeValues;
+      for (auto staticSize : staticSizes)
+        staticSizeValues.push_back(rewriter.create<ConstantIntOp>(
+            loc, rewriter.getI64IntegerAttr(staticSize)));
+      Value argList = rewriter.create<PrimListConstructOp>(
+          loc, Torch::ListType::get(rewriter.getType<Torch::IntType>()),
+          staticSizeValues);
+      rewriter.replaceOpWithNewOp<AtenViewOp>(op, op.getType(), op.self(),
+                                              argList);
+      return success();
+    }
+    return failure();
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // AtenDimOp
 //===----------------------------------------------------------------------===//

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

@@ -464,8 +464,8 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::to.other : (Tensor, Tensor, bool, bool, int?) -> (Tensor)")
     emit("aten::to.prim_Device : (Tensor, Device?, int?, bool, bool) -> (Tensor)")
     emit("aten::to.device : (Tensor, Device, int, bool, bool, int?) -> (Tensor)")
-    emit("aten::type_as : (Tensor, Tensor) -> (Tensor)")
-    emit("aten::view : (Tensor, int[]) -> (Tensor)", has_folder=True)
+    emit("aten::type_as : (Tensor, Tensor) -> (Tensor)", has_folder=True)
+    emit("aten::view : (Tensor, int[]) -> (Tensor)", has_folder=True, has_canonicalizer=True)
     emit("aten::_unsafe_view : (Tensor, int[]) -> (Tensor)")
     emit("aten::where.self : (Tensor, Tensor, Tensor) -> (Tensor)")
     emit("aten::where.Scalar : (Tensor, Scalar, Scalar) -> (Tensor)")

python/torch_mlir_e2e_test/test_suite/constant_alloc.py

@@ -1462,3 +1462,26 @@ 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=1).to(torch.float)
+    module.forward(tu.randint(2, 3, low=-10, high=10),
+                   mask)
+

python/torch_mlir_e2e_test/test_suite/reshape_like.py

@@ -259,6 +259,25 @@ def View1DFoldModule_basic(module, tu: TestUtils):
 
 # ==============================================================================
 
+class ViewStaticModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([5, 12, 3, 8], torch.float32, True),
+    ])
+
+    def forward(self, a):
+        return a.view(6, 5, 6, 8)
+
+@register_test_case(module_factory=lambda: ViewStaticModule())
+def ViewStaticModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(5, 12, 3, 8))
+
+# ==============================================================================
+
 class ViewCollapseInferredDimModule(torch.nn.Module):
     def __init__(self):
         super().__init__()

python/torch_mlir_e2e_test/test_suite/type_conversion.py

@@ -214,3 +214,24 @@ def forward(self, x):
 @register_test_case(module_factory=lambda: ToDtypeBoolLayoutNoneModule())
 def ToDtypeBoolLayoutNoneModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 5))
+
+class TypeAsSameModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+        ([-1, -1], torch.float32, True),
+    ])
+    def forward(self, x, y):
+        return torch.ops.aten.type_as(x,
+                                 y)
+
+
+
+@register_test_case(module_factory=lambda: TypeAsSameModule())
+def TypeAsSameModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 5), tu.rand(3, 5))