Repository: apache/arrow
File: cpp/src/arrow/gpu/cuda_test.cc (C++)

Code before:
  AssertCudaBufferEquals(checked_cast<const CudaBuffer&>(buffer),
                         std::forward<Expected>(expected));
}

class TestCudaBase : public ::testing::Test {
 public:
  void SetUp() {
    ASSERT_OK_AND_ASSIGN(manager_, CudaDeviceManager::Instance());
    ASSERT_OK_AND_ASSIGN(device_, manager_->GetDevice(kGpuNumber));
    //     ASSERT_OK(device_->GetContext(kGpuNumber, &context_));
    ASSERT_OK_AND_ASSIGN(context_, device_->GetContext());
    ASSERT_OK_AND_ASSIGN(mm_, AsCudaMemoryManager(device_->default_memory_manager()));
    cpu_device_ = CPUDevice::Instance();
    cpu_mm_ = cpu_device_->default_memory_manager();
  }

  void TearDown() {
    for (auto cu_context : non_primary_contexts_) {
      ASSERT_CUDA_OK(cuCtxDestroy(cu_context));
    }
  }

  Result<CUcontext> NonPrimaryRawContext() {
    CUcontext ctx;
#if CUDA_VERSION >= 13000
    RETURN_NOT_OK(StatusFromCuda(cuCtxCreate_v4(&ctx, /*ctxCreateParams=*/nullptr,
                                                /*flags=*/0, device_->handle())));
#else
    RETURN_NOT_OK(StatusFromCuda(cuCtxCreate(&ctx, /*flags=*/0, device_->handle())));
#endif
    non_primary_contexts_.push_back(ctx);
    return ctx;
  }

  Result<std::shared_ptr<CudaContext>> NonPrimaryContext() {
    ARROW_ASSIGN_OR_RAISE(auto cuctx, NonPrimaryRawContext());
    return device_->GetSharedContext(cuctx);
  }

  // Returns nullptr if kOtherGpuNumber does not correspond to an installed GPU
  Result<std::shared_ptr<CudaDevice>> OtherGpuDevice() {
    auto maybe_device = CudaDevice::Make(kOtherGpuNumber);
    if (maybe_device.status().IsInvalid()) {
      return nullptr;
    }
    return maybe_device;
  }

 protected:
  CudaDeviceManager* manager_;
  std::shared_ptr<CudaDevice> device_;

Review comment (bug):
I'm not sure `cuCtxCreate_v4` is designated as a public API. For CUDA 13, `cuCtxCreate` is `#define`d as `cuCtxCreate_v4`, so I think I'd be inclined to write this as:

```suggestion
    RETURN_NOT_OK(StatusFromCuda(cuCtxCreate(&ctx, /*ctxCreateParams=*/nullptr,
```

Code after:
  AssertCudaBufferEquals(checked_cast<const CudaBuffer&>(buffer),
                         std::forward<Expected>(expected));
}

class TestCudaBase : public ::testing::Test {
 public:
  void SetUp() {
    ASSERT_OK_AND_ASSIGN(manager_, CudaDeviceManager::Instance());
    ASSERT_OK_AND_ASSIGN(device_, manager_->GetDevice(kGpuNumber));
    //     ASSERT_OK(device_->GetContext(kGpuNumber, &context_));
    ASSERT_OK_AND_ASSIGN(context_, device_->GetContext());
    ASSERT_OK_AND_ASSIGN(mm_, AsCudaMemoryManager(device_->default_memory_manager()));
    cpu_device_ = CPUDevice::Instance();
    cpu_mm_ = cpu_device_->default_memory_manager();
  }

  void TearDown() {
    for (auto cu_context : non_primary_contexts_) {
      ASSERT_CUDA_OK(cuCtxDestroy(cu_context));
    }
  }

  Result<CUcontext> NonPrimaryRawContext() {
    CUcontext ctx;
#if CUDA_VERSION >= 13000
    RETURN_NOT_OK(StatusFromCuda(cuCtxCreate(&ctx, /*ctxCreateParams=*/nullptr,
                                             /*flags=*/0, device_->handle())));
#else
    RETURN_NOT_OK(StatusFromCuda(cuCtxCreate(&ctx, /*flags=*/0, device_->handle())));
#endif
    non_primary_contexts_.push_back(ctx);
    return ctx;
  }

  Result<std::shared_ptr<CudaContext>> NonPrimaryContext() {
    ARROW_ASSIGN_OR_RAISE(auto cuctx, NonPrimaryRawContext());
    return device_->GetSharedContext(cuctx);
  }

  // Returns nullptr if kOtherGpuNumber does not correspond to an installed GPU
  Result<std::shared_ptr<CudaDevice>> OtherGpuDevice() {
    auto maybe_device = CudaDevice::Make(kOtherGpuNumber);
    if (maybe_device.status().IsInvalid()) {
      return nullptr;
    }
    return maybe_device;
  }

 protected:
  CudaDeviceManager* manager_;
  std::shared_ptr<CudaDevice> device_;