[Proto] Import CPU memory for Taichi AOT

c_api/include/taichi/taichi_cpu.h

@@ -20,6 +20,10 @@ ti_export_cpu_memory(TiRuntime runtime,
                      TiMemory memory,
                      TiCpuMemoryInteropInfo *interop_info);
 
+TI_DLL_EXPORT TiMemory TI_API_CALL ti_import_cpu_memory(TiRuntime runtime,
+                                                        void *ptr,
+                                                        size_t memory_size);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif  // __cplusplus

c_api/src/taichi_llvm_impl.cpp

@@ -167,6 +167,24 @@ void ti_export_cpu_memory(TiRuntime runtime,
 #endif  // TI_WITH_LLVM
 }
 
+TI_DLL_EXPORT TiMemory TI_API_CALL ti_import_cpu_memory(TiRuntime runtime,
+                                                        void *ptr,
+                                                        size_t memory_size) {
+  capi::LlvmRuntime *llvm_runtime =
+      static_cast<capi::LlvmRuntime *>((Runtime *)runtime);
+
+  auto &device = llvm_runtime->get();
+  auto &cpu_device = static_cast<taichi::lang::cpu::CpuDevice &>(device);
+
+  taichi::lang::DeviceAllocation device_alloc =
+      cpu_device.import_memory(ptr, memory_size);
+
+  // prepare memory object
+  TiMemory memory = devalloc2devmem(*llvm_runtime, device_alloc);
+
+  return memory;
+}
+
 // function.export_cuda_runtime
 void ti_export_cuda_memory(TiRuntime runtime,
                            TiMemory memory,

c_api/tests/c_api_behavior_test.cpp

@@ -4,6 +4,7 @@
 #include "gtest/gtest.h"
 #include "c_api_test_utils.h"
 #include "taichi/cpp/taichi.hpp"
+#include "taichi/taichi_cpu.h"
 #include "c_api/tests/gtest_fixture.h"
 
 TEST_F(CapiTest, TestBehaviorCreateRuntime) {

c_api/tests/c_api_interop_test.cpp

@@ -86,4 +86,37 @@ TEST_F(CapiTest, AotTestVulkanTextureInterop) {
   }
 }
 
+TEST_F(CapiTest, TestCPUImport) {
+  TiArch arch = TiArch::TI_ARCH_X64;
+  ti::Runtime runtime(arch);
+
+  float data_x[4] = {1.0, 2.0, 3.0, 4.0};
+
+  auto memory = ti_import_cpu_memory(runtime, &data_x[0], sizeof(float) * 4);
+
+  int dim_count = 1;
+  int element_count = 4;
+  auto elem_type = TI_DATA_TYPE_F32;
+
+  // prepare tiNdArray
+  TiNdArray tiNdArray;
+  tiNdArray.memory = memory;
+  tiNdArray.shape.dim_count = dim_count;
+  tiNdArray.shape.dims[0] = element_count;
+  tiNdArray.elem_shape.dim_count = 0;
+  tiNdArray.elem_type = elem_type;
+
+  auto ti_memory = ti::Memory(runtime, memory, sizeof(float) * 4, false);
+  // prepare ndarray
+  auto ndarray = ti::NdArray<float>(std::move(ti_memory), tiNdArray);
+
+  std::vector<float> data_out(4);
+  ndarray.read(data_out);
+
+  std::cout << data_out[0] << std::endl;
+  std::cout << data_out[1] << std::endl;
+  std::cout << data_out[2] << std::endl;
+  std::cout << data_out[3] << std::endl;
+}
+
 #endif  // TI_WITH_VULKAN

tests/python/test_ipython.ipynb

@@ -1,5 +1,5 @@
 {
- "cells": [
+  "cells": [
   {
    "cell_type": "markdown",
    "id": "51b3b00e",