Loading...
Searching...
No Matches
CodeGen_PyTorch.h
Go to the documentation of this file.
#ifndef HALIDE_CODEGEN_PYTORCH_H
#define HALIDE_CODEGEN_PYTORCH_H
/** \file
*
* Defines an IRPrinter that emits C++ code that:
* 1. wraps PyTorch's C++ tensor into Halide * buffers,
* 2. calls the corresponding Halide operator.
* 3. maps the output buffer back to a PyTorch tensor.
*
* The generated code checks for runtime errors and raises PyTorch exception
* accordingly. It also makes sure the GPU device and stream are consistent when
* the PyTorch input, when applicable.
*/
#include "IRPrinter.h"
namespace Halide {
class Module;
namespace Internal {
struct LoweredFunc;
/** This class emits C++ code to wrap a Halide pipeline so that it can
* be used as a C++ extension operator in PyTorch.
*/
class CodeGen_PyTorch : public IRPrinter {
public:
CodeGen_PyTorch(std::ostream &dest);
~CodeGen_PyTorch() override = default;
/** Emit the PyTorch C++ wrapper for the Halide pipeline. */
void compile(const Module &module);
private:
void compile(const LoweredFunc &func, bool is_cuda);
};
} // namespace Internal
} // namespace Halide
#endif // HALIDE_CODEGEN_PYTORCH_H
IRPrinter.h
This header file defines operators that let you dump a Halide expression, statement,...
Halide::Internal::CodeGen_PyTorch::CodeGen_PyTorch
CodeGen_PyTorch(std::ostream &dest)
Halide::Internal::CodeGen_PyTorch::~CodeGen_PyTorch
~CodeGen_PyTorch() override=default
Halide::Internal::CodeGen_PyTorch::compile
void compile(const Module &module)
Emit the PyTorch C++ wrapper for the Halide pipeline.
Halide::Internal::IRPrinter::IRPrinter
IRPrinter(std::ostream &)
Construct an IRPrinter pointed at a given output stream (e.g.
Halide::Module
A halide module.
Definition Module.h:142
Halide::Internal
Definition AbstractGenerator.h:24
Halide
This file defines the class FunctionDAG, which is our representation of a Halide pipeline,...
Definition AbstractGenerator.h:19
Halide::Internal::LoweredFunc
Definition of a lowered function.
Definition Module.h:101