From 26e56abbf8a6ac16ec1c4350ab8ac1a87c12387e Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 3 Sep 2023 20:12:30 -0400 Subject: [PATCH 1/2] check status of allocate_temp Fix #2513. Signed-off-by: Jinzhe Zeng --- source/op/prod_env_mat_multi_device.cc | 45 ++++++++++++++++++++------ 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/source/op/prod_env_mat_multi_device.cc b/source/op/prod_env_mat_multi_device.cc index a3acf8ea66..6472f015e2 100644 --- a/source/op/prod_env_mat_multi_device.cc +++ b/source/op/prod_env_mat_multi_device.cc @@ -1501,8 +1501,11 @@ static int _norm_copy_coord_gpu(OpKernelContext* context, // Tensor FPTYPE_temp; TensorShape FPTYPE_shape; FPTYPE_shape.AddDim(nall * 3); - context->allocate_temp(DataTypeToEnum::value, FPTYPE_shape, - tensor_list); + tensorflow::Status status = context->allocate_temp( + DataTypeToEnum::value, FPTYPE_shape, tensor_list); + if (!status.ok()) { + return false; + } FPTYPE* tmp_coord = (*tensor_list).flat().data(); DPErrcheck(cudaMemcpy(tmp_coord, coord, sizeof(FPTYPE) * nall * 3, cudaMemcpyDeviceToDevice)); @@ -1519,8 +1522,11 @@ static int _norm_copy_coord_gpu(OpKernelContext* context, // Tensor double_temp; TensorShape double_shape; double_shape.AddDim(18); - context->allocate_temp(DataTypeToEnum::value, double_shape, - tensor_list + 1); + tensorflow::Status status = context->allocate_temp( + DataTypeToEnum::value, double_shape, tensor_list + 1); + if (!status.ok()) { + return false; + } // Tensor int_temp; TensorShape int_shape; int_shape.AddDim(23 + nloc * 3 + loc_cellnum + total_cellnum * 3 + @@ -1583,7 +1589,11 @@ static int _build_nlist_gpu(OpKernelContext* context, // Tensor nlist_temp; TensorShape nlist_shape; nlist_shape.AddDim(nloc * 2); - context->allocate_temp(DT_INT32, nlist_shape, tensor_list); + tensorflow::Status status = + context->allocate_temp(DT_INT32, nlist_shape, tensor_list); + if (!status.ok()) { + return false; + } ilist = (*tensor_list).flat().data(); numneigh = ilist + nloc; // Tensor jlist_temp; @@ -1594,7 +1604,11 @@ static int _build_nlist_gpu(OpKernelContext* context, for (tt = 0; tt < max_nnei_trial; ++tt) { TensorShape jlist_shape; jlist_shape.AddDim(3 * int_64(nloc) * mem_nnei); - context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); + tensorflow::Status status = + context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); + if (!status.ok()) { + return false; + } jlist = (*(tensor_list + 1)).flat().data(); ind_data = jlist + nloc * mem_nnei; for (int_64 ii = 0; ii < nloc; ++ii) { @@ -1742,8 +1756,11 @@ static int _norm_copy_coord_gpu_rocm(OpKernelContext* context, // Tensor double_temp; TensorShape double_shape; double_shape.AddDim(18); - context->allocate_temp(DataTypeToEnum::value, double_shape, - tensor_list + 1); + tensorflow::Status status = context->allocate_temp( + DataTypeToEnum::value, double_shape, tensor_list + 1); + if (!status.ok()) { + return false; + } // Tensor int_temp; TensorShape int_shape; int_shape.AddDim(23 + nloc * 3 + loc_cellnum + total_cellnum * 3 + @@ -1806,7 +1823,11 @@ static int _build_nlist_gpu_rocm(OpKernelContext* context, // Tensor nlist_temp; TensorShape nlist_shape; nlist_shape.AddDim(nloc * 2); - context->allocate_temp(DT_INT32, nlist_shape, tensor_list); + tensorflow::Status status = + context->allocate_temp(DT_INT32, nlist_shape, tensor_list); + if (!status.ok()) { + return false; + } ilist = (*tensor_list).flat().data(); numneigh = ilist + nloc; // Tensor jlist_temp; @@ -1817,7 +1838,11 @@ static int _build_nlist_gpu_rocm(OpKernelContext* context, for (tt = 0; tt < max_nnei_trial; ++tt) { TensorShape jlist_shape; jlist_shape.AddDim(3 * int_64(nloc) * mem_nnei); - context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); + tensorflow::Status status = + context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); + if (!status.ok()) { + return false; + } jlist = (*(tensor_list + 1)).flat().data(); ind_data = jlist + nloc * mem_nnei; for (int_64 ii = 0; ii < nloc; ++ii) { From 3d691ff296b125932ad65668c7cc13bb3d24d83d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 3 Sep 2023 20:24:14 -0400 Subject: [PATCH 2/2] fix compile error Signed-off-by: Jinzhe Zeng --- source/op/prod_env_mat_multi_device.cc | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/source/op/prod_env_mat_multi_device.cc b/source/op/prod_env_mat_multi_device.cc index 6472f015e2..d423e8a108 100644 --- a/source/op/prod_env_mat_multi_device.cc +++ b/source/op/prod_env_mat_multi_device.cc @@ -1522,8 +1522,8 @@ static int _norm_copy_coord_gpu(OpKernelContext* context, // Tensor double_temp; TensorShape double_shape; double_shape.AddDim(18); - tensorflow::Status status = context->allocate_temp( - DataTypeToEnum::value, double_shape, tensor_list + 1); + status = context->allocate_temp(DataTypeToEnum::value, double_shape, + tensor_list + 1); if (!status.ok()) { return false; } @@ -1604,8 +1604,7 @@ static int _build_nlist_gpu(OpKernelContext* context, for (tt = 0; tt < max_nnei_trial; ++tt) { TensorShape jlist_shape; jlist_shape.AddDim(3 * int_64(nloc) * mem_nnei); - tensorflow::Status status = - context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); + status = context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); if (!status.ok()) { return false; } @@ -1838,8 +1837,7 @@ static int _build_nlist_gpu_rocm(OpKernelContext* context, for (tt = 0; tt < max_nnei_trial; ++tt) { TensorShape jlist_shape; jlist_shape.AddDim(3 * int_64(nloc) * mem_nnei); - tensorflow::Status status = - context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); + status = context->allocate_temp(DT_INT32, jlist_shape, tensor_list + 1); if (!status.ok()) { return false; }