Fix tests

Author: echuraev

src/relay/transforms/device_domains.cc

@@ -236,12 +236,13 @@ DeviceDomainPtr DeviceDomains::DomainForCallee(const Call& call) {
     args_and_result.emplace_back(ForVirtualDevice(device_copy_props.body->checked_type(),
                                                   device_copy_props.dst_virtual_device));
   } else if (call->op == alloc_storage_op) {
-    ICHECK_EQ(call->args.size(), 2U);
-    // alloc_storage(size, alignment, virtual_device=<t>)
-    // alloc_storage: fn(<cpu>, <cpu>):<t>
+    ICHECK_EQ(call->args.size(), 3U);
+    // alloc_storage(size, shape, alignment, virtual_device=<t>)
+    // alloc_storage: fn(<cpu>, <cpu>, <cpu>):<t>
     const auto* attrs = call->attrs.as<AllocStorageAttrs>();
     args_and_result.emplace_back(host_domain_);
     args_and_result.emplace_back(host_domain_);
+    args_and_result.emplace_back(host_domain_);
     args_and_result.emplace_back(ForVirtualDevice(call->checked_type(), attrs->virtual_device));
   } else if (call->op == alloc_tensor_op) {
     ICHECK_EQ(call->args.size(), 3U);

tests/python/relay/test_pass_dead_code_elimination.py

@@ -16,8 +16,10 @@
 # under the License.
 import tvm
 import tvm.testing
+from tvm import relay
 from tvm.relay import Function, transform
 from tvm.relay.testing import inception_v3
+import numpy as np
 import pytest
 
 cpu_scope = tvm.target.VirtualDevice(tvm.cpu(), tvm.target.Target("llvm"))
@@ -228,14 +230,19 @@ def @main() {
 
 
 def test_impure_op():
+    shape = np.array([64, 2])
+    metatable = {
+        "VirtualDevice": [cpu_scope],
+        "relay.Constant": [relay.const(shape, dtype="int64")],
+    }
     """Don't elide calls to side-effecting operators."""
     before_program = tvm.relay.parse(
         """
         #[version = "0.0.5"]
         def @main() {
            let %size: int64 = cast(1024, dtype="int64");
            let %alignment: int64 = cast(64, dtype="int64");
-           let %x = memory.alloc_storage(%size, %alignment, virtual_device=meta[VirtualDevice][0]);
+           let %x = memory.alloc_storage(%size, meta[relay.Constant][0], %alignment, virtual_device=meta[VirtualDevice][0]);
            let %_ = memory.kill(%x);
            0
         }
@@ -250,6 +257,7 @@ def @main() {
         #[version = "0.0.5"]
         def @main() {
            %0 = memory.alloc_storage(cast(1024, dtype="int64"),
+                                     meta[relay.Constant][0],
                                      cast(64, dtype="int64"),
                                      virtual_device=meta[VirtualDevice][0]);
            let %_ = memory.kill(%0);
@@ -267,14 +275,19 @@ def @main() {
 
 
 def test_impure_func():
+    shape = np.array([64, 2])
+    metatable = {
+        "VirtualDevice": [cpu_scope],
+        "relay.Constant": [relay.const(shape, dtype="int64")],
+    }
     """Don't elide calls to side-effecting functions."""
     before_program = tvm.relay.parse(
         """
         #[version = "0.0.5"]
         def @f() -> int {
            let %size: int64 = cast(1024, dtype="int64");
            let %alignment: int64 = cast(64, dtype="int64");
-           let %x = memory.alloc_storage(%size, %alignment, virtual_device=meta[VirtualDevice][0]);
+           let %x = memory.alloc_storage(%size, meta[relay.Constant][0], %alignment, virtual_device=meta[VirtualDevice][0]);
            let %_ = memory.kill(%x);
            0
         }
@@ -293,6 +306,7 @@ def @main() -> int {
         #[version = "0.0.5"]
         def @f() -> int {
            %0 = memory.alloc_storage(cast(1024, dtype="int64"),
+                                     meta[relay.Constant][0],
                                      cast(64, dtype="int64"),
                                      virtual_device=meta[VirtualDevice][0]);
            let %_ = memory.kill(%0);

tests/python/relay/test_pass_plan_devices.py

@@ -761,14 +761,18 @@ def ref(x):
 
 
 def test_alloc_storage():
-    metatable = {"VirtualDevice": [HOST, GPU]}
+    shape = np.array([3, 2])
+    metatable = {
+        "VirtualDevice": [HOST, GPU],
+        "relay.Constant": [relay.const(shape, dtype="int64")],
+    }
 
     def input():
         return tvm.relay.parse(
             """
             #[version = "0.0.5"]
             def @main(%size: int64, %alignment: int64) {
-              memory.alloc_storage(%size, %alignment, virtual_device=meta[VirtualDevice][1])
+              memory.alloc_storage(%size, meta[relay.Constant][0], %alignment, virtual_device=meta[VirtualDevice][1])
             }
         """,
             "from_string",
@@ -782,7 +786,8 @@ def expected():
             #[version = "0.0.5"]
             def @main(%size {virtual_device=meta[VirtualDevice][0]}: int64, %alignment {virtual_device=meta[VirtualDevice][0]}: int64,
                       virtual_device=meta[VirtualDevice][1]) {
-              memory.alloc_storage(%size, %alignment, virtual_device=meta[VirtualDevice][1])
+              %0 = on_device(meta[relay.Constant][0], virtual_device=meta[VirtualDevice][0], constrain_result=True);
+              memory.alloc_storage(%size, %0, %alignment, virtual_device=meta[VirtualDevice][1])
             }
         """,
             "from_string",