[TORCH][MLIR] Add E2E support for `aten.[ones_like|zeros_like]`

- This commit adds E2E support for `aten.ones_like` and
  `aten.zeros_like` ops.
- Adds support for non-None `dtype` argument of `aten.empty_like` op.
- All the unit test cases related to constant tensor allocation like ops
  are moved to a different file named `constant_alloc.py`.

Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>

e2e_testing/torchscript/basic.py

@@ -586,144 +586,6 @@ def ExpandModule_basic(module, tu: TestUtils):
 
 # ==============================================================================
 
-
-class OnesModuleInt(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.ones(3, 4, dtype=torch.int64)
-
-@register_test_case(module_factory=lambda: OnesModuleInt())
-def OnesModuleInt_basic(module, tu: TestUtils):
-    module.forward()
-
-# ==============================================================================
-
-class OnesModuleFloat(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.ones(3, 4, dtype=torch.float32)
-
-@register_test_case(module_factory=lambda: OnesModuleFloat())
-def OnesModuleFloat_basic(module, tu: TestUtils):
-    module.forward()
-
-
-class OnesModuleFalsePinMemory(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.ones(3, 4, dtype=torch.float32, pin_memory=False)
-
-@register_test_case(module_factory=lambda: OnesModuleFalsePinMemory())
-def OnesModuleFalsePinMemory_basic(module, tu: TestUtils):
-    module.forward()
-
-# ==============================================================================
-
-class EmptyIntModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return 0 * torch.empty((3, 4), dtype=torch.int64)
-
-@register_test_case(module_factory=lambda: EmptyIntModule())
-def EmptyModule_int(module, tu: TestUtils):
-    module.forward()
-
-# ==============================================================================
-
-class EmptyFloatModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.pow(torch.empty((3, 4), dtype=torch.float32), 0)
-
-@register_test_case(module_factory=lambda: EmptyFloatModule())
-def EmptyModule_float(module, tu: TestUtils):
-    module.forward()
-
-
-class EmptyFalsePinMemoryModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.pow(torch.empty((3, 4), dtype=torch.float32, 
-                                     pin_memory=False), 0)
-
-@register_test_case(module_factory=lambda: EmptyFalsePinMemoryModule())
-def EmptyModule_falsePinMemory(module, tu: TestUtils):
-    module.forward()
-
-# ==============================================================================
-
-class EmptyLikeIntModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-        ([-1, -1], torch.int64, True),
-    ])
-    def forward(self, a):
-        return 0 * torch.empty_like(a, dtype=torch.int64)
-
-@register_test_case(module_factory=lambda: EmptyLikeIntModule())
-def EmptyLikeModule_int(module, tu: TestUtils):
-    module.forward(torch.randint(10, (3, 5)))
-
-# ==============================================================================
-
-class EmptyLikeFloatModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-        ([-1, -1], torch.float32, True),
-    ])
-    def forward(self, a):
-        return torch.pow(torch.empty_like(a, dtype=torch.float32), 0)
-
-@register_test_case(module_factory=lambda: EmptyLikeFloatModule())
-def EmptyLikeModule_float(module, tu: TestUtils):
-    module.forward(tu.rand(4, 5))
-
-# ==============================================================================
-
 class ContiguousModule(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -926,57 +788,6 @@ def DropoutModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
 
 
-class Fill_TensorFloat64WithFloat32(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-        ([-1, -1, -1], torch.float32, True),
-    ])
-    def forward(self, tensor):
-        return torch.ops.aten.fill_(tensor, 3.0)
-
-@register_test_case(module_factory=lambda: Fill_TensorFloat64WithFloat32())
-def Fill_TensorFloat64WithFloat32_basic(module, tu: TestUtils):
-    module.forward(torch.randn(3, 2, 4))
-
-
-class Fill_TensorFloat64WithFloat64(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-        ([-1, -1, -1], torch.float64, True),
-    ])
-    def forward(self, tensor):
-        return torch.ops.aten.fill_(tensor, 3.0)
-
-@register_test_case(module_factory=lambda: Fill_TensorFloat64WithFloat64())
-def Fill_TensorFloat64WithFloat64_basic(module, tu: TestUtils):
-    module.forward(torch.randn(3, 2, 4).to(torch.float64))
-
-
-class Fill_TensorFloat64WithInt64(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-        ([-1, -1, -1], torch.float64, True),
-    ])
-    def forward(self, tensor):
-        return torch.ops.aten.fill_(tensor, 3)
-
-@register_test_case(module_factory=lambda: Fill_TensorFloat64WithInt64())
-def Fill_TensorFloat64WithInt64_basic(module, tu: TestUtils):
-    module.forward(torch.randn(3, 2, 4).to(torch.float64))
-
-
 class MeanModule(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -1047,86 +858,6 @@ def NumelZeroRankModule_basic(module, tu: TestUtils):
     module.forward(torch.randint(10,[]))
 
 
-class ZerosModuleInt2D(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.zeros(3, 4, dtype=torch.int64)
-
-@register_test_case(module_factory=lambda: ZerosModuleInt2D())
-def ZerosModuleInt2D_basic(module, tu: TestUtils):
-    module.forward()
-
-
-class ZerosModuleInt3D(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.zeros(3, 4, 5, dtype=torch.int64)
-
-@register_test_case(module_factory=lambda: ZerosModuleInt3D())
-def ZerosModuleInt3D_basic(module, tu: TestUtils):
-    module.forward()
-
-
-class ZerosModuleFloat2D(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.zeros(3, 4, dtype=torch.float32)
-
-@register_test_case(module_factory=lambda: ZerosModuleFloat2D())
-def ZerosModuleFloat2D_basic(module, tu: TestUtils):
-    module.forward()
-
-
-class ZerosModuleFloat3D(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.zeros(3, 4, 5, dtype=torch.float32)
-
-@register_test_case(module_factory=lambda: ZerosModuleFloat3D())
-def ZerosModuleFloat3D_basic(module, tu: TestUtils):
-    module.forward()
-
-
-class ZerosModuleFalsePinMemory(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    @export
-    @annotate_args([
-        None,
-    ])
-    def forward(self):
-        return torch.zeros(3, 4, dtype=torch.float32, pin_memory=False)
-
-@register_test_case(module_factory=lambda: ZerosModuleFalsePinMemory())
-def ZerosModuleFalsePinMemory_basic(module, tu: TestUtils):
-    module.forward()
-
-
 class BoolTensorReturnFalseModule(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -1181,6 +912,7 @@ def BoolTensorReturnMixedModule_basic(module, tu: TestUtils):
     module.forward(torch.tensor([[1, 0], [0,1]], dtype=torch.bool))
 
 # ==============================================================================
+
 class TModuleRank2(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -1193,11 +925,11 @@ class TModuleRank2(torch.nn.Module):
     def forward(self, lhs):
         return torch.t(lhs)
 
-
 @register_test_case(module_factory=lambda: TModuleRank2())
 def TModuleRank2_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
 
+
 class TModuleRank1(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -1210,11 +942,11 @@ class TModuleRank1(torch.nn.Module):
     def forward(self, lhs):
         return torch.t(lhs)
 
-
 @register_test_case(module_factory=lambda: TModuleRank1())
 def TModuleRank1_basic(module, tu: TestUtils):
     module.forward(tu.rand(3))
 
+
 class TModuleRank0(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -1227,7 +959,6 @@ class TModuleRank0(torch.nn.Module):
     def forward(self, lhs):
         return torch.t(lhs)
 
-
 @register_test_case(module_factory=lambda: TModuleRank0())
 def TModuleRank0_basic(module, tu: TestUtils):
     module.forward(torch.tensor(7, dtype=torch.float32))

e2e_testing/torchscript/constant_alloc.py

@@ -0,0 +1,400 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# Also available under a BSD-style license. See LICENSE.
+
+import torch
+
+from torch_mlir_e2e_test.torchscript.framework import TestUtils
+from torch_mlir_e2e_test.torchscript.registry import register_test_case
+from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+
+# ==============================================================================
+
+class ZerosModuleInt2D(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.zeros(3, 4, dtype=torch.int64)
+
+@register_test_case(module_factory=lambda: ZerosModuleInt2D())
+def ZerosModuleInt2D_basic(module, tu: TestUtils):
+    module.forward()
+
+
+class ZerosModuleInt3D(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.zeros(3, 4, 5, dtype=torch.int64)
+
+@register_test_case(module_factory=lambda: ZerosModuleInt3D())
+def ZerosModuleInt3D_basic(module, tu: TestUtils):
+    module.forward()
+
+
+class ZerosModuleFloat2D(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.zeros(3, 4, dtype=torch.float32)
+
+@register_test_case(module_factory=lambda: ZerosModuleFloat2D())
+def ZerosModuleFloat2D_basic(module, tu: TestUtils):
+    module.forward()
+
+
+class ZerosModuleFloat3D(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.zeros(3, 4, 5, dtype=torch.float32)
+
+@register_test_case(module_factory=lambda: ZerosModuleFloat3D())
+def ZerosModuleFloat3D_basic(module, tu: TestUtils):
+    module.forward()
+
+
+class ZerosModuleFalsePinMemory(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.zeros(3, 4, dtype=torch.float32, pin_memory=False)
+
+@register_test_case(module_factory=lambda: ZerosModuleFalsePinMemory())
+def ZerosModuleFalsePinMemory_basic(module, tu: TestUtils):
+    module.forward()
+
+# ==============================================================================
+
+class OnesModuleInt(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.ones(3, 4, dtype=torch.int64)
+
+@register_test_case(module_factory=lambda: OnesModuleInt())
+def OnesModuleInt_basic(module, tu: TestUtils):
+    module.forward()
+
+
+class OnesModuleFloat(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.ones(3, 4, dtype=torch.float32)
+
+@register_test_case(module_factory=lambda: OnesModuleFloat())
+def OnesModuleFloat_basic(module, tu: TestUtils):
+    module.forward()
+
+
+class OnesModuleFalsePinMemory(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.ones(3, 4, dtype=torch.float32, pin_memory=False)
+
+@register_test_case(module_factory=lambda: OnesModuleFalsePinMemory())
+def OnesModuleFalsePinMemory_basic(module, tu: TestUtils):
+    module.forward()
+
+# ==============================================================================
+
+class EmptyIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return 0 * torch.empty((3, 4), dtype=torch.int64)
+
+@register_test_case(module_factory=lambda: EmptyIntModule())
+def EmptyModule_int(module, tu: TestUtils):
+    module.forward()
+
+
+class EmptyFloatModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.pow(torch.empty((3, 4), dtype=torch.float32), 0)
+
+@register_test_case(module_factory=lambda: EmptyFloatModule())
+def EmptyModule_float(module, tu: TestUtils):
+    module.forward()
+
+
+class EmptyFalsePinMemoryModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.pow(torch.empty((3, 4), dtype=torch.float32, 
+                                     pin_memory=False), 0)
+
+@register_test_case(module_factory=lambda: EmptyFalsePinMemoryModule())
+def EmptyModule_falsePinMemory(module, tu: TestUtils):
+    module.forward()
+
+# ==============================================================================
+
+class EmptyLikeIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.int64, True),
+    ])
+    def forward(self, a):
+        return 0 * torch.empty_like(a, dtype=torch.int32)
+
+@register_test_case(module_factory=lambda: EmptyLikeIntModule())
+def EmptyLikeModule_int(module, tu: TestUtils):
+    module.forward(torch.randint(10, (3, 5)))
+
+
+class EmptyLikeFloatModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.pow(torch.empty_like(a, dtype=torch.float32), 0)
+
+@register_test_case(module_factory=lambda: EmptyLikeFloatModule())
+def EmptyLikeModule_float(module, tu: TestUtils):
+    module.forward(tu.rand(4, 5))
+
+
+class EmptyLikeFalsePinMemoryModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.pow(torch.empty_like(a, dtype=torch.float64, 
+                                     pin_memory=False), 0)
+
+@register_test_case(module_factory=lambda: EmptyLikeFalsePinMemoryModule())
+def EmptyLikeModule_falsePinMemory(module, tu: TestUtils):
+    module.forward(tu.rand(2, 3, 4))
+
+# ==============================================================================
+
+class ZerosLikeIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.int64, True),
+    ])
+    def forward(self, a):
+        return torch.zeros_like(a, dtype=torch.int32)
+
+@register_test_case(module_factory=lambda: ZerosLikeIntModule())
+def ZerosLikeModule_int(module, tu: TestUtils):
+    module.forward(torch.randint(10, (3, 5)))
+
+
+class ZerosLikeFloatModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.zeros_like(a, dtype=torch.float32)
+
+@register_test_case(module_factory=lambda: ZerosLikeFloatModule())
+def ZerosLikeModule_float(module, tu: TestUtils):
+    module.forward(tu.rand(4, 5))
+
+
+class ZerosLikeFalsePinMemoryModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.zeros_like(a, dtype=torch.float64, pin_memory=False)
+
+@register_test_case(module_factory=lambda: ZerosLikeFalsePinMemoryModule())
+def ZerosLikeModule_falsePinMemory(module, tu: TestUtils):
+    module.forward(tu.rand(2, 3, 4))
+
+# ==============================================================================
+
+class OnesLikeIntModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.int64, True),
+    ])
+    def forward(self, a):
+        return torch.ones_like(a, dtype=torch.int32)
+
+@register_test_case(module_factory=lambda: OnesLikeIntModule())
+def OnesLikeModule_int(module, tu: TestUtils):
+    module.forward(torch.randint(10, (3, 5)))
+
+
+class OnesLikeFloatModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.ones_like(a, dtype=torch.float32)
+
+@register_test_case(module_factory=lambda: OnesLikeFloatModule())
+def OnesLikeModule_float(module, tu: TestUtils):
+    module.forward(tu.rand(4, 5))
+
+
+class OnesLikeFalsePinMemoryModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float32, True),
+    ])
+    def forward(self, a):
+        return torch.ones_like(a, dtype=torch.float64, pin_memory=False)
+
+@register_test_case(module_factory=lambda: OnesLikeFalsePinMemoryModule())
+def OnesLikeModule_falsePinMemory(module, tu: TestUtils):
+    module.forward(tu.rand(2, 3, 4))
+
+# ==============================================================================
+
+class Fill_TensorFloat64WithFloat32(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float32, True),
+    ])
+    def forward(self, tensor):
+        return torch.ops.aten.fill_(tensor, 3.0)
+
+@register_test_case(module_factory=lambda: Fill_TensorFloat64WithFloat32())
+def Fill_TensorFloat64WithFloat32_basic(module, tu: TestUtils):
+    module.forward(torch.randn(3, 2, 4))
+
+
+class Fill_TensorFloat64WithFloat64(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float64, True),
+    ])
+    def forward(self, tensor):
+        return torch.ops.aten.fill_(tensor, 3.0)
+
+@register_test_case(module_factory=lambda: Fill_TensorFloat64WithFloat64())
+def Fill_TensorFloat64WithFloat64_basic(module, tu: TestUtils):
+    module.forward(torch.randn(3, 2, 4).to(torch.float64))
+
+
+class Fill_TensorFloat64WithInt64(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float64, True),
+    ])
+    def forward(self, tensor):
+        return torch.ops.aten.fill_(tensor, 3)
+
+@register_test_case(module_factory=lambda: Fill_TensorFloat64WithInt64())
+def Fill_TensorFloat64WithInt64_basic(module, tu: TestUtils):
+    module.forward(torch.randn(3, 2, 4).to(torch.float64))
+

e2e_testing/torchscript/main.py

@@ -46,6 +46,7 @@ from . import slice_like
 from . import nll_loss
 from . import index_select
 from . import arange
+from . import constant_alloc
 
 def _get_argparse():
     config_choices = ['native_torch', 'torchscript', 'refbackend', 'tosa', 'external']

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

@@ -2160,6 +2160,44 @@ def Torch_AtenEmptyLikeOp : Torch_Op<"aten.empty_like", [
   let assemblyFormat = "$self `,` $dtype `,` $layout `,` $device `,` $pin_memory `,` $memory_format attr-dict `:` type($self) `,` type($dtype) `,` type($layout) `,` type($device) `,` type($pin_memory) `,` type($memory_format) `->` type($result)";
 }
 
+def Torch_AtenZerosLikeOp : Torch_Op<"aten.zeros_like", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::zeros_like : (Tensor, int?, int?, Device?, bool?, int?) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    TorchOptionalIntType:$dtype,
+    TorchOptionalIntType:$layout,
+    TorchOptionalDeviceType:$device,
+    TorchOptionalBoolType:$pin_memory,
+    TorchOptionalIntType:$memory_format
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self `,` $dtype `,` $layout `,` $device `,` $pin_memory `,` $memory_format attr-dict `:` type($self) `,` type($dtype) `,` type($layout) `,` type($device) `,` type($pin_memory) `,` type($memory_format) `->` type($result)";
+}
+
+def Torch_AtenOnesLikeOp : Torch_Op<"aten.ones_like", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::ones_like : (Tensor, int?, int?, Device?, bool?, int?) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    TorchOptionalIntType:$dtype,
+    TorchOptionalIntType:$layout,
+    TorchOptionalDeviceType:$device,
+    TorchOptionalBoolType:$pin_memory,
+    TorchOptionalIntType:$memory_format
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self `,` $dtype `,` $layout `,` $device `,` $pin_memory `,` $memory_format attr-dict `:` type($self) `,` type($dtype) `,` type($layout) `,` type($device) `,` type($pin_memory) `,` type($memory_format) `->` type($result)";
+}
+
 def Torch_AtenEmptyMemoryFormatOp : Torch_Op<"aten.empty.memory_format", [
     AllowsTypeRefinement,
     HasValueSemantics

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -536,8 +536,8 @@ class DecomposeAtenLayerNormOp : public OpRewritePattern<AtenLayerNormOp> {
 };
 } // namespace
 
-// Decompose `aten.empty_like` op into `aten.size` and `aten.empty` ops.
 namespace {
+// Decompose `aten.empty_like` op into `aten.size` and `aten.empty` ops.
 class DecomposeAtenEmptyLikeOp : public OpRewritePattern<AtenEmptyLikeOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
@@ -547,14 +547,6 @@ public:
         Torch::ListType::get(Torch::IntType::get(op.getContext()));
     Value sizeList =
         rewriter.create<AtenSizeOp>(op.getLoc(), sizeListType, op.self());
-
-    // TODO: Handle the case when `dtype` is NoneType.
-    Type dtype = op.dtype().getType();
-    if (dtype.isa<OptionalType>() || dtype.isa<Torch::NoneType>() ||
-        dtype.isa<mlir::NoneType>())
-      return rewriter.notifyMatchFailure(
-          op, "unimplemented: None dtype is not supported");
-
     rewriter.replaceOpWithNewOp<AtenEmptyMemoryFormatOp>(
         op, op.getType(), sizeList, op.dtype(), op.layout(), op.device(),
         op.pin_memory(), op.memory_format());
@@ -605,6 +597,28 @@ class DecomposeAtenArangeStartOp : public OpRewritePattern<AtenArangeStartOp> {
 };
 } // namespace
 
+namespace {
+// Decompose constant tensor allocation like ops.
+template <typename OpTy, int fillVal>
+class DecomposeConstantTensorAllocLikeOp : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    // Allocate a memory block.
+    Value initTensor = rewriter.create<AtenEmptyLikeOp>(
+        loc, op.getType(), op.self(), op.dtype(), op.layout(), op.device(),
+        op.pin_memory(), op.memory_format());
+    Value constVal = rewriter.create<Torch::ConstantIntOp>(
+        loc, rewriter.getI64IntegerAttr(fillVal));
+    // Initialize the allocated memory block with `fillVal`.
+    rewriter.replaceOpWithNewOp<AtenFill_ScalarOp>(op, initTensor.getType(),
+                                                   initTensor, constVal);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class DecomposeComplexOpsPass
     : public DecomposeComplexOpsBase<DecomposeComplexOpsPass> {
@@ -622,6 +636,12 @@ class DecomposeComplexOpsPass
     target.addIllegalOp<AtenLogSoftmaxIntOp>();
     patterns.add<DecomposeAtenEmptyLikeOp>(context);
     target.addIllegalOp<AtenEmptyLikeOp>();
+    patterns.add<DecomposeConstantTensorAllocLikeOp<AtenOnesLikeOp, 1>>(
+        context);
+    target.addIllegalOp<AtenOnesLikeOp>();
+    patterns.add<DecomposeConstantTensorAllocLikeOp<AtenZerosLikeOp, 0>>(
+        context);
+    target.addIllegalOp<AtenZerosLikeOp>();
     patterns.add<DecomposeAtenExpandOp>(context);
     target.addIllegalOp<AtenExpandOp>();
     patterns.add<DecomposeAtenSizeOp>(context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -234,15 +234,15 @@ public:
   visitOperation(Operation *op,
                  ArrayRef<LatticeElement<ValueKnowledge> *> operands) final {
     if (isa<TensorStaticInfoCastOp, CopyToValueTensorOp, CopyToNonValueTensorOp,
-            AtenTanhOp, AtenBatchNormOp, AtenReluOp, AtenGeluOp,
+            AtenTanhOp, AtenBatchNormOp, AtenReluOp, AtenGeluOp, AtenCeilOp,
             AtenGeluBackwardOp, AtenBitwiseNotOp, AtenExpOp, AtenSinOp,
             AtenCosOp, AtenSigmoidOp, DerefineOp, AtenToPrimDeviceOp, AtenCpuOp,
-            AtenContiguousOp, AtenFill_ScalarOp, AtenDetachOp, AtenEmptyLikeOp,
+            AtenContiguousOp, AtenFill_ScalarOp, AtenDetachOp, AtenReciprocalOp,
             AtenMaskedFill_ScalarOp, AtenCopy_Op, AtenCumsumOp, AtenLayerNormOp,
             AtenClampOp, AtenLogOp, AtenNegOp, AtenSqrtOp, AtenFloorOp,
             AtenLog2Op, Aten_SoftmaxBackwardDataOp, AtenRsqrtOp, AtenDropoutOp,
             AtenTanhBackwardOp, Aten_LogSoftmaxBackwardDataOp, AtenAddIntOp,
-            AtenAbsOp, AtenReciprocalOp, AtenCeilOp>(op)) {
+            AtenAbsOp>(op)) {
       return getLatticeElement(op->getResult(0)).join(*operands[0]);
     }
 
@@ -395,6 +395,14 @@ public:
     } else if (auto emptyMemoryFormat = dyn_cast<AtenEmptyMemoryFormatOp>(op)) {
       return visitConstantTensorAllocOp<AtenEmptyMemoryFormatOp>(
           emptyMemoryFormat);
+    } else if (auto zerosLike = dyn_cast<AtenZerosLikeOp>(op)) {
+      return visitConstantTensorAllocLikeOp<AtenZerosLikeOp>(zerosLike,
+                                                             operands);
+    } else if (auto onesLike = dyn_cast<AtenOnesLikeOp>(op)) {
+      return visitConstantTensorAllocLikeOp<AtenOnesLikeOp>(onesLike, operands);
+    } else if (auto emptyLike = dyn_cast<AtenEmptyLikeOp>(op)) {
+      return visitConstantTensorAllocLikeOp<AtenEmptyLikeOp>(emptyLike,
+                                                             operands);
     } else if (auto toDtype = dyn_cast<AtenToDtypeOp>(op)) {
       return visitAtenToDtypeOp(toDtype, operands);
     } else if (auto toOther = dyn_cast<AtenToOtherOp>(op)) {
@@ -566,6 +574,9 @@ private:
   ChangeResult visitScalarToTensorConversionOp(OpTy op);
   ChangeResult visitAtenTensorOp(AtenTensorOp op);
   template <typename OpTy> ChangeResult visitConstantTensorAllocOp(OpTy op);
+  template <typename OpTy>
+  ChangeResult visitConstantTensorAllocLikeOp(
+      OpTy op, ArrayRef<LatticeElement<ValueKnowledge> *> operands);
   ChangeResult
   visitAtenToDtypeOp(AtenToDtypeOp op,
                      ArrayRef<LatticeElement<ValueKnowledge> *> operands);
@@ -1407,6 +1418,20 @@ ChangeResult TypeAnalyzer::visitConstantTensorAllocOp(OpTy op) {
   return getLatticeElement(op.getResult()).join(knowledge);
 }
 
+template <typename OpTy>
+ChangeResult TypeAnalyzer::visitConstantTensorAllocLikeOp(
+    OpTy op, ArrayRef<LatticeElement<ValueKnowledge> *> operands) {
+  auto input = operands[0]->getValue();
+  auto knowledge =
+      ValueKnowledge::getNotNonePessimisticValueState(op->getContext());
+  if (input.hasSizes) {
+    knowledge.hasSizes = true;
+    knowledge.sizes = input.sizes;
+  }
+  fillInDTypeGivenDTypeAndDataType(knowledge, op.dtype(), input.dtype);
+  return getLatticeElement(op.getResult()).join(knowledge);
+}
+
 // Convert input tensor type to the given `dtype`.
 ChangeResult TypeAnalyzer::visitAtenToDtypeOp(
     AtenToDtypeOp op, ArrayRef<LatticeElement<ValueKnowledge> *> operands) {

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

@@ -565,6 +565,8 @@ def emit_aten_ops(torch_ir_dir: str, registry: Registry):
         emit("aten::detach : (Tensor) -> (Tensor)")
         emit("aten::embedding : (Tensor, Tensor, int, bool, bool) -> (Tensor)")
         emit("aten::empty_like : (Tensor, int?, int?, Device?, bool?, int?) -> (Tensor)")
+        emit("aten::zeros_like : (Tensor, int?, int?, Device?, bool?, int?) -> (Tensor)")
+        emit("aten::ones_like : (Tensor, int?, int?, Device?, bool?, int?) -> (Tensor)")
         emit("aten::empty.memory_format : (int[], int?, int?, Device?, bool?, int?) -> (Tensor)")
         emit("aten::expand : (Tensor, int[], bool) -> (Tensor)")
         emit("aten::broadcast_to : (Tensor, int[]) -> (Tensor)")