Loading...
Searching...
No Matches
HalidePyTorchCudaHelpers.h
Go to the documentation of this file.
#ifndef HL_PYTORCH_CUDA_HELPERS_H
#define HL_PYTORCH_CUDA_HELPERS_H
/** \file
* Override Halide's CUDA hooks so that the Halide code called from PyTorch uses
* the correct GPU device and stream. This header should be included once in
* the PyTorch/C++ binding source file (see apps/HelloPyTorch/setup.py for an
* example).
*/
#include "HalideRuntimeCuda.h"
#include "cuda.h"
#include "cuda_runtime.h"
namespace Halide {
namespace PyTorch {
typedef struct UserContext {
UserContext(int id, CUcontext *ctx, cudaStream_t *stream)
: device_id(id), cuda_context(ctx), stream(stream) {};
int device_id;
CUcontext *cuda_context;
cudaStream_t *stream;
} UserContext;
} // namespace PyTorch
} // namespace Halide
// Replace Halide weakly-linked CUDA handles
extern "C" {
int halide_cuda_acquire_context(void *user_context, CUcontext *ctx, bool create = true) {
if (user_context != nullptr) {
Halide::PyTorch::UserContext *user_ctx = (Halide::PyTorch::UserContext *)user_context;
*ctx = *user_ctx->cuda_context;
} else {
*ctx = nullptr;
}
return halide_error_code_success;
}
int halide_cuda_get_stream(void *user_context, CUcontext ctx, CUstream *stream) {
if (user_context != nullptr) {
Halide::PyTorch::UserContext *user_ctx = (Halide::PyTorch::UserContext *)user_context;
*stream = *user_ctx->stream;
} else {
*stream = 0;
}
return halide_error_code_success;
}
int halide_get_gpu_device(void *user_context) {
if (user_context != nullptr) {
Halide::PyTorch::UserContext *user_ctx = (Halide::PyTorch::UserContext *)user_context;
return user_ctx->device_id;
} else {
return 0;
}
}
} // extern "C"
#endif /* end of include guard: HL_PYTORCH_CUDA_HELPERS_H */
halide_cuda_acquire_context
int halide_cuda_acquire_context(void *user_context, CUcontext *ctx, bool create=true)
Definition HalidePyTorchCudaHelpers.h:33
halide_cuda_get_stream
int halide_cuda_get_stream(void *user_context, CUcontext ctx, CUstream *stream)
Definition HalidePyTorchCudaHelpers.h:43
halide_get_gpu_device
int halide_get_gpu_device(void *user_context)
Halide calls this to get the desired halide gpu device setting.
Definition HalidePyTorchCudaHelpers.h:53
halide_error_code_success
@ halide_error_code_success
There was no error.
Definition HalideRuntime.h:1076
HalideRuntimeCuda.h
Routines specific to the Halide Cuda runtime.
Halide::PyTorch
Definition HalidePyTorchCudaHelpers.h:16
Halide::PyTorch::UserContext
struct Halide::PyTorch::UserContext UserContext
Halide
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...
Definition AbstractGenerator.h:19
Halide::PyTorch::UserContext
Definition HalidePyTorchCudaHelpers.h:18
Halide::PyTorch::UserContext::cuda_context
CUcontext * cuda_context
Definition HalidePyTorchCudaHelpers.h:23
Halide::PyTorch::UserContext::device_id
int device_id
Definition HalidePyTorchCudaHelpers.h:22
Halide::PyTorch::UserContext::UserContext
UserContext(int id, CUcontext *ctx, cudaStream_t *stream)
Definition HalidePyTorchCudaHelpers.h:19
Halide::PyTorch::UserContext::stream
cudaStream_t * stream
Definition HalidePyTorchCudaHelpers.h:24