From 8a3834835e18c0cddf1cb2c51c1d14cf65fe888b Mon Sep 17 00:00:00 2001
From: Mark Shields <mbs@octoml.ai>
Date: Wed, 10 Nov 2021 16:30:24 -0800
Subject: [PATCH] [checkpoint] Fix after rebase on CallLowered.

---
 src/relay/op/memory/device_copy.cc     | 13 -------------
 src/relay/transforms/device_domains.cc |  6 ++----
 2 files changed, 2 insertions(+), 17 deletions(-)

diff --git a/src/relay/op/memory/device_copy.cc b/src/relay/op/memory/device_copy.cc
index 538264ce9688c..48d12368fa285 100644
--- a/src/relay/op/memory/device_copy.cc
+++ b/src/relay/op/memory/device_copy.cc
@@ -105,19 +105,6 @@ DeviceCopyProps GetDeviceCopyProps(const CallNode* call_node) {
     } else {
       return {call_node->args[0], device_copy_attrs->src_se_scope, device_copy_attrs->dst_se_scope};
     }
-  } else if (call_node->op == CallLoweredOp()) {
-    /* Get device props for a TIR function */
-    CallLoweredProps call_lowered_props = GetCallLoweredProps(call_node);
-
-    if (call_lowered_props.attrs.metadata.count("source_device") == 1 &&
-        call_lowered_props.attrs.metadata.count("dst_device") == 1) {
-      ICHECK_EQ(call_lowered_props.arguments.size(), 1) << "device_copy is of arity 1";
-      return {call_lowered_props.lowered_func,
-              static_cast<DLDeviceType>(
-                  Downcast<Integer>(call_lowered_props.attrs.metadata["source_device"])->value),
-              static_cast<DLDeviceType>(
-                  Downcast<Integer>(call_lowered_props.attrs.metadata["dst_device"])->value)};
-    }
   }
   return {};
 }
diff --git a/src/relay/transforms/device_domains.cc b/src/relay/transforms/device_domains.cc
index 667379d7a9a0c..44c0ecf41de8e 100644
--- a/src/relay/transforms/device_domains.cc
+++ b/src/relay/transforms/device_domains.cc
@@ -51,10 +51,8 @@ DeviceCopyProps GetPrimitiveDeviceCopyProps(const CallNode* call_node) {
         call_lowered_props.attrs.metadata.count("dst_device") == 1) {
       ICHECK_EQ(call_lowered_props.arguments.size(), 1) << "device_copy is of arity 1";
       return {call_lowered_props.arguments[0],
-              static_cast<DLDeviceType>(
-                  Downcast<Integer>(call_lowered_props.attrs.metadata["source_device"])->value),
-              static_cast<DLDeviceType>(
-                  Downcast<Integer>(call_lowered_props.attrs.metadata["dst_device"])->value)};
+              Downcast<SEScope>(call_lowered_props.attrs.metadata["src_se_scope"]),
+              Downcast<SEScope>(call_lowered_props.attrs.metadata["dst_se_scope"])};
     }
   }
   return {};