modified workspace return

Author: xgqdut2016

src/ops/random_sample/cuda/random_sample.cu

@@ -109,8 +109,8 @@ void inclusive_sum(
         stream);
 }
 
-void random_sample_workspace(size_t &size_radix_sort, size_t &size_scan,
-                             int voc, DT dtype) {
+infiniopStatus_t random_sample_workspace(size_t &size_radix_sort, size_t &size_scan,
+                                         int voc, DT dtype) {
     if (dtype_eq(dtype, F16)) {
         sort_pairs_descending<half, uint64_t>(nullptr, size_radix_sort,
                                               nullptr, nullptr,
@@ -121,6 +121,7 @@ void random_sample_workspace(size_t &size_radix_sort, size_t &size_scan,
             nullptr, size_scan,
             nullptr, voc,
             nullptr);
+        return STATUS_SUCCESS;
     } else if (dtype_eq(dtype, F32)) {
         sort_pairs_descending<float, uint64_t>(nullptr, size_radix_sort,
                                                nullptr, nullptr,
@@ -131,6 +132,7 @@ void random_sample_workspace(size_t &size_radix_sort, size_t &size_scan,
             nullptr, size_scan,
             nullptr, voc,
             nullptr);
+        return STATUS_SUCCESS;
     } else if (dtype_eq(dtype, F64)) {
         sort_pairs_descending<double, uint64_t>(nullptr, size_radix_sort,
                                                 nullptr, nullptr,
@@ -141,8 +143,9 @@ void random_sample_workspace(size_t &size_radix_sort, size_t &size_scan,
             nullptr, size_scan,
             nullptr, voc,
             nullptr);
+        return STATUS_SUCCESS;
     } else {
-        throw std::invalid_argument("Unsupported dtype provided.");
+        return STATUS_BAD_TENSOR_DTYPE;
     }
 }
 __global__ void random_sample_kernel(uint64_t *result,

src/ops/random_sample/cuda/random_sample.cuh

@@ -19,8 +19,8 @@ typedef struct RandomSampleCudaDescriptor *RandomSampleCudaDescriptor_t;
 infiniopStatus_t cudaCreateRandomSampleDescriptor(CudaHandle_t handle,
                                                   RandomSampleCudaDescriptor_t *desc_ptr, infiniopTensorDescriptor_t result,
                                                   infiniopTensorDescriptor_t probs);
-void random_sample_workspace(size_t &size_radix_sort, size_t &size_scan,
-                             int voc, DT dtype);
+infiniopStatus_t random_sample_workspace(size_t &size_radix_sort, size_t &size_scan,
+                                         int voc, DT dtype);
 infiniopStatus_t cudaGetRandomSampleWorkspaceSize(RandomSampleCudaDescriptor_t desc, unsigned long int *size);
 
 infiniopStatus_t cudaRandomSample(RandomSampleCudaDescriptor_t desc,

src/ops/random_sample/cuda/random_sample_cuda.cc

@@ -33,8 +33,11 @@ infiniopStatus_t cudaCreateRandomSampleDescriptor(CudaHandle_t handle,
 infiniopStatus_t cudaGetRandomSampleWorkspaceSize(RandomSampleCudaDescriptor_t desc, unsigned long int *size) {
     size_t size_radix_sort;
     size_t size_scan;
-    random_sample_workspace(size_radix_sort, size_scan,
-                            desc->voc, desc->dtype);
+    infiniopStatus_t status = random_sample_workspace(size_radix_sort, size_scan,
+                                                      desc->voc, desc->dtype);
+    if (status != STATUS_SUCCESS) {
+        return status;
+    }
     *size = desc->step + std::max(size_radix_sort, size_scan);
     return STATUS_SUCCESS;
 }