Loading...
Searching...
No Matches
HalidePyTorchHelpers.h
Go to the documentation of this file.
#ifndef HL_PYTORCH_WRAPPER_H
#define HL_PYTORCH_WRAPPER_H
/** \file
* Set of utility functions to wrap PyTorch tensors into Halide buffers,
* making sure the data in on the correct device (CPU/GPU). This header
* is included in each generated op by the PyTorch CodeGen.
*/
#include <exception>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
#include "HalideBuffer.h"
// Forward declare the cuda_device_interface, for tensor wrapper.
extern "C" const halide_device_interface_t *halide_cuda_device_interface();
#define HLPT_CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define HLPT_CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
#define HLPT_CHECK_DEVICE(x, dev) AT_ASSERTM(x.is_cuda() && x.get_device() == dev, #x " must be a CUDA tensor")
namespace Halide {
namespace PyTorch {
using Halide::Runtime::Buffer;
inline std::vector<int> get_dims(const at::Tensor tensor) {
int ndims = tensor.ndimension();
std::vector<int> dims(ndims, 0);
// PyTorch dim order is reverse of Halide
for (int dim = 0; dim < ndims; ++dim) {
dims[dim] = tensor.size(ndims - 1 - dim);
}
return dims;
}
template<class scalar_t>
inline void check_type(at::Tensor &tensor) {
AT_ERROR("Scalar type ", tensor.scalar_type(), " not handled by Halide's PyTorch wrapper");
}
// TODO: if PyTorch exposes any variable with the API version,
// I haven't found it in source or documentation; for now, we'll sniff
// this macro's existence to infer that we are building with v1.3+ (vs 1.2)
#ifdef AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS
#define HL_PYTORCH_API_VERSION 13
#else
#define HL_PYTORCH_API_VERSION 12
#endif
#if HL_PYTORCH_API_VERSION >= 13
// PyTorch 1.3+
#define HL_PT_DEFINE_TYPECHECK(ctype, ttype) \
template<> \
inline void check_type<ctype>(at::Tensor & tensor) { \
AT_ASSERTM(tensor.scalar_type() == at::ScalarType::ttype, "scalar type do not match"); \
}
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS(HL_PT_DEFINE_TYPECHECK);
#undef HL_PT_DEFINE_TYPECHECK
#else // HL_PYTORCH_API_VERSION < 13
// PyTorch 1.2
#define HL_PT_DEFINE_TYPECHECK(ctype, ttype, _3) \
template<> \
inline void check_type<ctype>(at::Tensor & tensor) { \
AT_ASSERTM(tensor.scalar_type() == at::ScalarType::ttype, "scalar type do not match"); \
}
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(HL_PT_DEFINE_TYPECHECK);
#undef HL_PT_DEFINE_TYPECHECK
#endif // HL_PYTORCH_API_VERSION check
template<class scalar_t>
inline Buffer<scalar_t> wrap(at::Tensor &tensor) {
check_type<scalar_t>(tensor);
std::vector<int> dims = get_dims(tensor);
#if HL_PYTORCH_API_VERSION >= 13
scalar_t *pData = tensor.data_ptr<scalar_t>();
#else
scalar_t *pData = tensor.data<scalar_t>();
#endif
return Buffer<scalar_t>(pData, dims);
}
template<class scalar_t>
inline Buffer<scalar_t> wrap_cuda(at::Tensor &tensor) {
check_type<scalar_t>(tensor);
std::vector<int> dims = get_dims(tensor);
#if HL_PYTORCH_API_VERSION >= 13
scalar_t *pData = tensor.data_ptr<scalar_t>();
#else
scalar_t *pData = tensor.data<scalar_t>();
#endif
AT_ASSERTM(tensor.is_cuda(), "expected input tensor to be on a CUDA device.");
Buffer<scalar_t> buffer(dims);
const halide_device_interface_t *cuda_interface = halide_cuda_device_interface();
int err = buffer.device_wrap_native(cuda_interface, (uint64_t)pData);
AT_ASSERTM(err == 0, "(CUDA) halide_device_wrap failed");
buffer.set_device_dirty();
return buffer;
}
} // namespace PyTorch
} // namespace Halide
#endif // HL_PYTORCH_WRAPPER_H
HalideBuffer.h
Defines a Buffer type that wraps from halide_buffer_t and adds functionality, and methods for more co...
halide_cuda_device_interface
const halide_device_interface_t * halide_cuda_device_interface()
HL_PT_DEFINE_TYPECHECK
#define HL_PT_DEFINE_TYPECHECK(ctype, ttype, _3)
Definition HalidePyTorchHelpers.h:71
Halide::Runtime::Buffer
A templated Buffer class that wraps halide_buffer_t and adds functionality.
Definition HalideBuffer.h:224
Halide::Runtime::Buffer::set_device_dirty
void set_device_dirty(bool v=true)
Definition HalideBuffer.h:1808
Halide::Runtime::Buffer::device_wrap_native
int device_wrap_native(const struct halide_device_interface_t *device_interface, uint64_t handle, void *ctx=nullptr)
Definition HalideBuffer.h:1854
Halide::PyTorch::get_dims
std::vector< int > get_dims(const at::Tensor tensor)
Definition HalidePyTorchHelpers.h:30
Halide::PyTorch::wrap
Buffer< scalar_t > wrap(at::Tensor &tensor)
Definition HalidePyTorchHelpers.h:84
Halide::PyTorch::wrap_cuda
Buffer< scalar_t > wrap_cuda(at::Tensor &tensor)
Definition HalidePyTorchHelpers.h:96
Halide::PyTorch::AT_FORALL_SCALAR_TYPES_WITH_COMPLEX
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(HL_PT_DEFINE_TYPECHECK)
Halide::PyTorch::check_type
void check_type(at::Tensor &tensor)
Definition HalidePyTorchHelpers.h:41
Halide
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...
Definition AbstractGenerator.h:19
uint64_t
unsigned __INT64_TYPE__ uint64_t
Definition runtime_internal.h:23
halide_device_interface_t
Each GPU API provides a halide_device_interface_t struct pointing to the code that manages device all...
Definition HalideRuntime.h:797