1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
|
HelloPytorch is a simple application that demonstrates how to build a PyTorch
wrapper around Halide ops using the internal Generator target.
The demo op takes two 4D tensors as input, `a` and `b` and outputs their
pointwise sum.
The application builds float32 and float64 versions of the op for CPU and CUDA
(provided that a GPU is present and the `nvcc` compiler is installed and in the
path, see `Makefile`).
The build proceeds in four steps:
1. Build the Halide generator
2. Build Halide code for each each datatype and each architecture
(`-e static_library,h` flag in the generator).
3. Build the PyTorch C++ interface code (`-e pytorch_wrapper` flag in the
generator). Note that the PyTorch wrapper requires the `user_context` feature
be enabled in the generator `target` for CUDA ops. This allows Halide's and
PyTorch's GPU memory managers to communicate with each other.
4. Synthesize and compile a Python extension that links togethers the various
operator libraries and exposes them to Python (see `setup.py`).
Building only requires Python 3 and PyTorch. Please follow these instructions to
install the latest PyTorch: https://pytorch.org/
If everything is setup correctly, running `make test` should build the PyTorch
extension and runs a simple test (`test.py`).
Tested with Python 3.7.4 and PyTorch 1.2.0
|
