mirror of
https://github.com/NVIDIA/nvbench.git
synced 2026-03-14 20:27:24 +00:00
78 lines
2.3 KiB
Plaintext
78 lines
2.3 KiB
Plaintext
/*
|
|
* Copyright 2023 NVIDIA Corporation
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 with the LLVM exception
|
|
* (the "License"); you may not use this file except in compliance with
|
|
* the License.
|
|
*
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://llvm.org/foundation/relicensing/LICENSE.txt
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include <nvbench/config.cuh>
|
|
#include <nvbench/cuda_stream.cuh>
|
|
#include <nvbench/device_manager.cuh>
|
|
#include <nvbench/types.cuh>
|
|
|
|
#include <fmt/format.h>
|
|
|
|
#include "test_asserts.cuh"
|
|
|
|
namespace
|
|
{
|
|
#ifdef NVBENCH_HAS_CUPTI
|
|
/**
|
|
* @brief Queries and returns the device id that the given \p cuda_stream is associated with
|
|
*
|
|
* @param cuda_stream The stream to get the device id for
|
|
* @return The device id that \p cuda_stream is associated with
|
|
*/
|
|
int get_device_of_stream(cudaStream_t cuda_stream)
|
|
{
|
|
CUcontext ctx;
|
|
NVBENCH_DRIVER_API_CALL(cuStreamGetCtx(CUstream{cuda_stream}, &ctx));
|
|
NVBENCH_DRIVER_API_CALL(cuCtxPushCurrent(ctx));
|
|
CUdevice device_id{};
|
|
NVBENCH_DRIVER_API_CALL(cuCtxGetDevice(&device_id));
|
|
NVBENCH_DRIVER_API_CALL(cuCtxPopCurrent(&ctx));
|
|
return static_cast<int>(device_id);
|
|
}
|
|
#endif
|
|
} // namespace
|
|
|
|
void test_basic()
|
|
{
|
|
#ifdef NVBENCH_HAS_CUPTI
|
|
// Get devices
|
|
auto devices = nvbench::device_manager::get().get_devices();
|
|
|
|
// Iterate over devices
|
|
for (auto const &device_info : devices)
|
|
{
|
|
// Create stream on the device before it becomes the active device
|
|
nvbench::cuda_stream device_stream(device_info);
|
|
|
|
// Verify cuda stream is associated with the correct cuda device
|
|
ASSERT(get_device_of_stream(device_stream.get_stream()) == device_info.get_id());
|
|
|
|
// Set the device as active device
|
|
device_info.set_active();
|
|
|
|
// Create the stream (implicitly) on the device that is currently active
|
|
nvbench::cuda_stream current_device_stream{};
|
|
|
|
// Verify the cuda stream was in fact associated with the currently active device
|
|
ASSERT(get_device_of_stream(current_device_stream.get_stream()) == device_info.get_id());
|
|
}
|
|
#endif
|
|
}
|
|
|
|
int main() { test_basic(); }
|