[TORCH][MLIR] Add E2E support for aten.clone (#571)

This commit adds support for the aten.clone op.

e2e_testing/torchscript/elementwise.py

@@ -1106,3 +1106,21 @@ class ElementwiseAddScalarFloatModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: ElementwiseAddScalarFloatModule())
 def ElementwiseAddScalarFloatModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
+
+
+class ElementwiseCloneModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float32, True),
+    ])
+    def forward(self, x):
+        return torch.clone(x)
+
+
+@register_test_case(module_factory=lambda: ElementwiseCloneModule())
+def ElementwiseCloneModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(2, 3, 4))

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

@@ -2262,6 +2262,21 @@ def Torch_AtenBucketizeTensorOp : Torch_Op<"aten.bucketize.Tensor", [
   let assemblyFormat = "$self `,` $boundaries `,` $out_int32 `,` $right attr-dict `:` qualified(type($self)) `,` qualified(type($boundaries)) `,` qualified(type($out_int32)) `,` qualified(type($right)) `->` qualified(type($result))";
 }
 
+def Torch_AtenCloneOp : Torch_Op<"aten.clone", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::clone : (Tensor, int?) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    TorchOptionalIntType:$memory_format
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self `,` $memory_format attr-dict `:` qualified(type($self)) `,` qualified(type($memory_format)) `->` qualified(type($result))";
+}
+
 def Torch_AtenContiguousOp : Torch_Op<"aten.contiguous", [
     AllowsTypeRefinement
   ]> {

lib/Conversion/TorchToLinalg/TorchToLinalg.cpp

@@ -1660,6 +1660,13 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
     return b.create<math::SqrtOp>(loc, payloadArgs[0]);
   if (isa<AtenRsqrtOp>(op))
     return b.create<math::RsqrtOp>(loc, payloadArgs[0]);
+  if (auto clone = dyn_cast<AtenCloneOp>(op)) {
+    if (!clone.memory_format().getType().isa<Torch::NoneType>()) {
+      clone.emitError("unimplemented: only default memory format is supported");
+      return nullptr;
+    }
+    return payloadArgs[0];
+  }
   if (auto bitwiseAndTensor = dyn_cast<AtenBitwiseAndTensorOp>(op)) {
     if (bitwiseAndTensor.getType()
             .cast<ValueTensorType>()
@@ -2450,7 +2457,7 @@ struct ConvertElementwiseOp : ConversionPattern {
              AtenBitwiseAndTensorOp, AtenGtScalarOp, AtenEqScalarOp,
              AtenLtScalarOp, AtenWhereSelfOp, AtenCeilOp, AtenGtTensorOp,
              AtenEqTensorOp, AtenLtTensorOp, AtenSubScalarOp, AtenAddScalarOp,
-             AtenThresholdOp, AtenThresholdBackwardOp>(op))
+             AtenThresholdOp, AtenThresholdBackwardOp, AtenCloneOp>(op))
       return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
 
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
@@ -4571,7 +4578,7 @@ public:
         AtenAbsOp, AtenReciprocalOp, AtenBitwiseAndTensorOp, AtenGtScalarOp,
         AtenEqScalarOp, AtenLtScalarOp, AtenWhereSelfOp, AtenGtTensorOp,
         AtenEqTensorOp, AtenLtTensorOp, AtenThresholdOp,
-        AtenThresholdBackwardOp>();
+        AtenThresholdBackwardOp, AtenCloneOp>();
     patterns.add<ConvertElementwiseOp>(typeConverter, context);
     target.addIllegalOp<AtenSqueezeOp>();
     patterns.add<ConvertAtenSqueezeOp>(typeConverter, context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -242,8 +242,8 @@ public:
             AtenClampOp, AtenLogOp, AtenNegOp, AtenSqrtOp, AtenFloorOp,
             AtenLog2Op, Aten_SoftmaxBackwardDataOp, AtenRsqrtOp, AtenDropoutOp,
             AtenTanhBackwardOp, Aten_LogSoftmaxBackwardDataOp, AtenAddIntOp,
-            AtenAbsOp, AtenThresholdOp, AtenSquareOp, PseudoAtenUniformOp>(
+            AtenAbsOp, AtenThresholdOp, AtenSquareOp, PseudoAtenUniformOp,
-            op)) {
+            AtenCloneOp>(op)) {
       return getLatticeElement(op->getResult(0)).join(*operands[0]);
     }
 

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

@@ -577,6 +577,7 @@ def emit_aten_ops(torch_ir_dir: str, registry: Registry):
         emit("aten::arange.start_step : (Scalar, Scalar, Scalar, int?, int?, Device?, bool?) -> (Tensor)")
         emit("aten::argmax : (Tensor, int?, bool) -> (Tensor)")
         emit("aten::bucketize.Tensor : (Tensor, Tensor, bool, bool) -> (Tensor)")
+        emit("aten::clone : (Tensor, int?) -> (Tensor)")
         emit("aten::contiguous : (Tensor, int) -> (Tensor)")
         emit("aten::copy_ : (Tensor, Tensor, bool) -> (Tensor)")
         emit("aten::detach : (Tensor) -> (Tensor)")