/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/jit/tensorexpr
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/jit/tensorexpr/llvm_codegen.h (3180B)
#pragma once
#ifdef TORCH_ENABLE_LLVM
#include
#include
#include
#include
#include
#include
#include
namespace torch {
namespace jit {
namespace tensorexpr {
class LLVMCodeGenImpl;
class TORCH_API LLVMCodeGen : public CodeGen {
public:
explicit LLVMCodeGen(
StmtPtr stmt,
const std::vector& args,
at::Device device = at::kCPU,
const std::string& kernel_func_name = "func",
Dtype dtype = kInt,
c10::optional triple = c10::nullopt,
c10::optional cpu = c10::nullopt,
c10::optional attrs = c10::nullopt);
explicit LLVMCodeGen(StmtPtr stmt);
LLVMCodeGen() = delete;
~LLVMCodeGen() override;
TORCH_API void call(const std::vector& args) override;
TORCH_API void call_raw(const std::vector& args) override;
at::Tensor empty_strided(
c10::IntArrayRef size,
c10::IntArrayRef stride,
c10::optional dtype_opt,
c10::optional layout_opt,
c10::optional device_opt,
c10::optional pin_memory_opt) override;
template
T value() {
return value(nullptr);
}
template
T value(std::vector& args) {
return value(args.data());
}
template
T value(void** args) {
T (*fp)(void**) = (T(*)(void**))getKernelAddress(impl_.get());
T rv = fp(args);
return rv;
}
std::string getCodeText(const std::string& attr = "") override;
private:
void* getKernelAddress(LLVMCodeGenImpl* impl);
std::unique_ptr impl_;
};
struct TORCH_API LLVMCodeGenBuilder {
using BufferArg = CodeGen::BufferArg;
LLVMCodeGenBuilder(StmtPtr stmt, std::vector args)
: stmt_(stmt), args_(std::move(args)) {}
LLVMCodeGenBuilder& device(at::Device device) {
device_ = device;
return *this;
}
LLVMCodeGenBuilder& kernelFuncName(std::string name) {
kernelFuncName_ = std::move(name);
return *this;
}
LLVMCodeGenBuilder& dtype(Dtype d) {
dtype_ = d;
return *this;
}
LLVMCodeGenBuilder& triple(std::string triple) {
triple_ = std::move(triple);
return *this;
}
LLVMCodeGenBuilder& cpu(std::string cpu) {
cpu_ = std::move(cpu);
return *this;
}
LLVMCodeGenBuilder& attrs(std::string attrs) {
attrs_ = std::move(attrs);
return *this;
}
std::unique_ptr build() {
return std::make_unique(
stmt_, args_, device_, kernelFuncName_, dtype_, triple_, cpu_, attrs_);
}
private:
StmtPtr stmt_;
std::vector args_;
at::Device device_ = at::kCPU;
std::string kernelFuncName_ = "func";
Dtype dtype_ = kInt;
c10::optional triple_ = c10::nullopt;
c10::optional cpu_ = c10::nullopt;
c10::optional attrs_ = c10::nullopt;
};
} // namespace tensorexpr
} // namespace jit
} // namespace torch
#endif // TORCH_ENABLE_LLVM