[Bugfix][Performance] Fix Column to copy 'self.index' to location of...

[Bugfix][Performance] Fix Column to copy 'self.index' to location of 'self.storage' only on read, and use kwargs. (#2267)

* Pass framework arguments to copy of frames

* Copy for subframe

* Fix mismatched contexts
Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com>

python/dgl/frame.py

@@ -115,7 +115,10 @@ class Column(object):
             # copy index to the same context of storage.
             # Copy index is usually cheaper than copy data
             if F.context(self.storage) != F.context(self.index):
-                self.index = F.copy_to(self.index, F.context(self.storage))
+                kwargs = {}
+                if self.device is not None:
+                    kwargs = self.device[1]
+                self.index = F.copy_to(self.index, F.context(self.storage), **kwargs)
             self.storage = F.gather_row(self.storage, self.index)
             self.index = None
 
@@ -148,8 +151,6 @@ class Column(object):
         """
         col = self.clone()
         col.device = (device, kwargs)
-        if self.index is not None:
-            col.index = F.copy_to(self.index, device)
         return col
 
     def __getitem__(self, rowids):
@@ -253,6 +254,12 @@ class Column(object):
         if self.index is None:
             return Column(self.storage, self.scheme, rowids, self.device)
         else:
+            if F.context(self.index) != F.context(rowids):
+                # make sure index and row ids are on the same context
+                kwargs = {}
+                if self.device is not None:
+                    kwargs = self.device[1]
+                rowids = F.copy_to(rowids, F.context(self.index), **kwargs)
             return Column(self.storage, self.scheme, F.gather_row(self.index, rowids), self.device)
 
     @staticmethod