1 #ifndef HL_PYTORCH_CUDA_HELPERS_H
2 #define HL_PYTORCH_CUDA_HELPERS_H
16 typedef struct UserContext {
17 UserContext(
int id,
CUcontext *ctx, cudaStream_t *stream)
18 : device_id(id), cuda_context(ctx), stream(stream){};
33 Halide::PyTorch::UserContext *user_ctx = (Halide::PyTorch::UserContext *)
user_context;
34 *ctx = *user_ctx->cuda_context;
43 Halide::PyTorch::UserContext *user_ctx = (Halide::PyTorch::UserContext *)
user_context;
44 *stream = *user_ctx->stream;
53 Halide::PyTorch::UserContext *user_ctx = (Halide::PyTorch::UserContext *)
user_context;
54 return user_ctx->device_id;
int halide_get_gpu_device(void *user_context)
Halide calls this to get the desired halide gpu device setting.
Routines specific to the Halide Cuda runtime.
struct CUstream_st * CUstream
CUDA stream.
struct CUctx_st * CUcontext
CUDA context.
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...