Skip to content

Commit

Permalink
Fixed address stability for NDArray args
Browse files Browse the repository at this point in the history
  • Loading branch information
PENGUINLIONG committed Jul 5, 2022
1 parent 56d5801 commit 22f5b97
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions c_api/src/taichi_core_impl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,8 @@ void ti_launch_compute_graph(TiRuntime runtime,

Runtime &runtime2 = *((Runtime *)runtime);
std::unordered_map<std::string, taichi::lang::aot::IValue> arg_map{};
std::vector<taichi::lang::Ndarray> ndarrays{};
std::vector<taichi::lang::Ndarray> ndarrays;
ndarrays.reserve(arg_count);

for (uint32_t i = 0; i < arg_count; ++i) {
const auto &arg = args[i];
Expand All @@ -304,12 +305,15 @@ void ti_launch_compute_graph(TiRuntime runtime,
return;
}
const TiNdArray &ndarray = arg.argument.value.ndarray;

std::vector<int> shape(ndarray.shape.dims,
ndarray.shape.dims + ndarray.shape.dim_count);

std::vector<int> elem_shape(ndarray.elem_shape.dims,
ndarray.elem_shape.dims + ndarray.elem_shape.dim_count);

ndarrays.emplace_back(taichi::lang::Ndarray(
devalloc, taichi::lang::PrimitiveType::f32, shape));
devalloc, taichi::lang::PrimitiveType::f32, shape, elem_shape));
arg_map.emplace(std::make_pair(
arg.name, taichi::lang::aot::IValue::create(ndarrays.back())));
break;
Expand Down

0 comments on commit 22f5b97

Please sign in to comment.