issue/591 - fix rearrange context mismatch

Author: wooway777

include/infinicore/context/context.hpp

@@ -17,6 +17,7 @@ size_t getDeviceCount(Device::Type type);
 
 infinirtStream_t getStream();
 infiniopHandle_t getInfiniopHandle();
+infiniopHandle_t getInfiniopHandle(Device device);
 
 void syncStream();
 void syncDevice();

src/infinicore/context/context_impl.cc

@@ -103,6 +103,13 @@ infiniopHandle_t getInfiniopHandle() {
     return ContextImpl::singleton().getCurrentRuntime()->infiniopHandle();
 }
 
+infiniopHandle_t getInfiniopHandle(Device device) {
+    if (device != getDevice()) {
+        setDevice(device);
+    }
+    return ContextImpl::singleton().getCurrentRuntime()->infiniopHandle();
+}
+
 void syncStream() {
     return ContextImpl::singleton().getCurrentRuntime()->syncStream();
 }

src/infinicore/ops/rearrange/rearrange_infiniop.cc

@@ -27,7 +27,7 @@ void calculate(Tensor y, Tensor x) {
     infiniopRearrangeDescriptor_t desc = nullptr;
 
     if (!desc_opt) {
-        INFINICORE_CHECK_ERROR(infiniopCreateRearrangeDescriptor(context::getInfiniopHandle(), &desc, y->desc(), x->desc()));
+        INFINICORE_CHECK_ERROR(infiniopCreateRearrangeDescriptor(context::getInfiniopHandle(y->device()), &desc, y->desc(), x->desc()));
         cache.put(seed, desc);
     } else {
         desc = *desc_opt;