Merge pull request #918 from InfiniTensor/issue/867

issue/867 fix cpu malloc

Author: PanZezhong1725

src/infinicore/context/context_impl.cc

@@ -126,6 +126,7 @@ std::shared_ptr<Memory> allocateMemory(size_t size) {
 }
 
 std::shared_ptr<Memory> allocateHostMemory(size_t size) {
+    setDevice(Device::cpu());
     return ContextImpl::singleton().getCpuRuntime()->allocateMemory(size);
 }
 

src/infiniop/ops/rearrange/nvidia/rearrange_nvidia.cu

@@ -345,12 +345,10 @@ infiniStatus_t launchKernel(
             const_cast<void *>(static_cast<const void *>(params.dst_grid_stride.data())),
             const_cast<void *>(static_cast<const void *>(constraints_data))};
 
-    CHECK_OR_RETURN(cudaLaunchKernel(
-                        kernel_func,
-                        static_cast<unsigned int>(grid_size), static_cast<unsigned int>(BLOCK_SIZE),
-                        args, 0, stream)
-                        == cudaSuccess,
-                    INFINI_STATUS_INTERNAL_ERROR);
+    CHECK_CUDA(cudaLaunchKernel(
+        kernel_func,
+        static_cast<unsigned int>(grid_size), static_cast<unsigned int>(BLOCK_SIZE),
+        args, 0, stream));
 
     return INFINI_STATUS_SUCCESS;
 }

xmake/test.lua

@@ -88,8 +88,6 @@ target("infinicore-test")
     add_files(os.projectdir().."/src/infinicore/ops/*/*.cc")
     add_files(os.projectdir().."/src/infinicore/nn/*.cc")
 
-    add_files(os.projectdir().."/src/infinicore-test/*.cc")
-    add_files(os.projectdir().."/src/infinicore-test/*/*.cc")
-
+    add_files(os.projectdir().."/src/infinicore-test/**.cc")
     set_installdir(INFINI_ROOT)
 target_end()