[GitHub] [tvm] areusch commented on a diff in pull request #12085: [Relay] Add RecoverVirtualDeviceMap helper

[GitBox <gi...@apache.org>]

areusch commented on code in PR #12085:
URL: https://github.com/apache/tvm/pull/12085#discussion_r920535891


##########
tests/python/relay/test_pass_plan_devices.py:
##########
@@ -241,6 +243,70 @@ def ref(a, b, c, d):
     exercise(input(), expected(), ref, rands((5, 7), 4))
 
 
+def test_left_add_on_cpu_via_copy_as_map():
+    metatable = {"VirtualDevice": [CPU, GPU]}
+
+    # As for test_left_add_on_cpu, but with an explicit device_copy.
+    def input():
+        return tvm.parser.parse(
+            """
+            #[version = "0.0.5"]
+            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
+                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32]) {
+              %0 = add(%a, %b);
+              %1 = device_copy(%0, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
+              %2 = add(%c, %d);
+              subtract(%1, %2)
+            }
+        """,
+            "from_string",
+            None,
+            metatable,
+        )
+
+    config = tvm.target.make_compilation_config(CTXT, TARGETS, HOST_TARGET)
+    actual_mod = relay.transform.InferType()(input())
+    actual_mod = relay.transform.PlanDevices(config)(actual_mod)
+    actual_mod = relay.transform.CapturePostDfsIndexInSpans()(actual_mod)
+
+    # Actual looks like:

Review Comment:
   does it make sense to assert this via assert str(actual_mod) == """..."""?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org