HalidePyTorchCudaHelpers.h File Reference

Override Halide's CUDA hooks so that the Halide code called from PyTorch uses the correct GPU device and stream. More...

#include "HalideRuntimeCuda.h"
#include "cuda.h"
#include "cuda_runtime.h"

Go to the source code of this file.

Classes
struct	Halide::PyTorch::UserContext

Namespaces
namespace	Halide
	This file defines the class FunctionDAG, which is our representation of a Halide pipeline, and contains methods to using Halide's bounds tools to query properties of it.
namespace	Halide::PyTorch

Typedefs
typedef struct Halide::PyTorch::UserContext	Halide::PyTorch::UserContext

Functions
int	halide_cuda_acquire_context (void *user_context, CUcontext *ctx, bool create=true)
int	halide_cuda_get_stream (void *user_context, CUcontext ctx, CUstream *stream)
int	halide_get_gpu_device (void *user_context)
	Halide calls this to get the desired halide gpu device setting.

Detailed Description

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).

Definition in file HalidePyTorchCudaHelpers.h.

Function Documentation

halide_cuda_acquire_context()

int halide_cuda_acquire_context	(	void *	user_context,
		CUcontext *	ctx,
		bool	create = true )

Definition at line 33 of file HalidePyTorchCudaHelpers.h.

References Halide::PyTorch::UserContext::cuda_context, and halide_error_code_success.

halide_cuda_get_stream()

int halide_cuda_get_stream	(	void *	user_context,
		CUcontext	ctx,
		CUstream *	stream )

Definition at line 43 of file HalidePyTorchCudaHelpers.h.

References halide_error_code_success, and Halide::PyTorch::UserContext::stream.

halide_get_gpu_device()

int halide_get_gpu_device

(

void *

user_context

)

Halide calls this to get the desired halide gpu device setting.

Implement this yourself to use a different gpu device per user_context. The default implementation returns the value set by halide_set_gpu_device, or the environment variable HL_GPU_DEVICE.

Definition at line 53 of file HalidePyTorchCudaHelpers.h.

References Halide::PyTorch::UserContext::device_id.