/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/kernel.h (9210B)
#pragma once
#include
#include
#include
#include
#include
#include
namespace torch {
namespace jit {
namespace tensorexpr {
// Returns true if the TE fuser supports this conv2d.
bool conv2dIsSupportedJit(const Node* node);
// Returns true if the TE fuser supports this matmul.
bool matmulIsSupported(const Node* node);
template
inline std::vector bufferSizes(const T& t) {
std::vector sizes;
for (size_t i = 0; i < t->ndim(); i++) {
sizes.push_back(*intValue(t->dim(i)));
}
return sizes;
}
enum ElementType {
kAllTypes = 0,
kIntegralTypes = 1 << 0,
kFloatingPointTypes = 1 << 1,
kBoolType = 1 << 2,
kComplexTypes = 1 << 3,
kQintTypes = 1 << 4,
kNonComplexOrQintTypes = kIntegralTypes | kBoolType | kFloatingPointTypes,
};
using ArgNone = c10::monostate;
using BufList = std::vector;
using IntList = std::vector;
using ArgValue = c10::variant<
tensorexpr::BufHandle,
tensorexpr::VarHandle,
double,
int64_t,
bool,
BufList,
IntList,
ArgNone>;
using NNCLoweringFunction = std::function&,
const std::vector&,
const c10::optional&,
at::Device)>;
// Get the dimensions of a value.
std::vector valueShape(const ArgValue& v);
// If v is a tensor, broadcast it to match the shape of axes, or return
// directly if v is a constant.
ExprHandle tensorOrConstant(
const ArgValue& v,
const std::vector& axes);
int64_t normalizeAndCheckIndex(int64_t idx, int64_t list_size);
ExprHandle broadcast(BufHandle b, const std::vector& axes);
ExprHandle constant(const ArgValue& v);
std::vector computeIndicesToBroadcast(
const std::vector& outputAxes,
const std::vector& inputSizes);
void promoteInputs(
std::vector& inputs,
const int typeConstraints = kAllTypes);
ExprHandle promoteToDtype(ExprHandle e, ScalarType dt);
ExprHandle promoteIntegerToDefaultType(const ExprHandle& e);
ExprHandle demoteOutput(
const ExprHandle& e,
const c10::optional type);
inline std::string getArgValueName(const ArgValue& a) {
if (c10::get_if(&a)) {
return "BufHandle";
} else if (c10::get_if(&a)) {
return "VarHandle";
} else if (c10::get_if(&a)) {
return "double";
} else if (c10::get_if(&a)) {
return "int64_t";
} else if (c10::get_if(&a)) {
return "bool";
} else if (c10::get_if(&a)) {
return "BufList";
} else if (c10::get_if(&a)) {
return "IntList";
} else if (c10::get_if(&a)) {
return "None";
} else {
throw std::runtime_error("ArgValue type not handled in string conversion");
}
}
template
std::vector convertVecArgValue(const std::vector& v) {
std::vector res;
for (auto& x : v) {
auto val = c10::get_if(&x);
if (val) {
res.push_back(*val);
} else {
throw std::runtime_error(
"vector type not homogeneous - found " + getArgValueName(x) +
", expected " + getArgValueName(v[0]));
}
}
return res;
}
struct TensorInfo {
std::vector dims;
c10::ScalarType dtype;
};
TORCH_API Tensor computeOperandValue(
c10::Symbol op,
const std::vector& inputs,
const std::vector& outputShape,
const c10::optional& outputType,
at::Device = at::kCPU);
class TORCH_API TensorExprKernel {
struct ConstantDescr {
BufPtr buf;
void* ptr;
};
public:
explicit TensorExprKernel(
const std::shared_ptr& subgraph,
std::unordered_map custom_lowerings =
{},
bool pre_alloc = false);
void run(Stack& stack);
void runFast(
const std::vector& inputs,
const std::vector& outputs);
void fallback(Stack& stack) {
InterpreterState(code_).run(stack);
}
StmtPtr getCodeGenStmt();
std::string getCodeText(const std::string& attr = "") {
return codegen_->getCodeText(attr);
}
const std::shared_ptr graph() {
return graph_;
}
const std::vector& getConstantDescriptors() const {
return constants_;
}
const std::vector& getBufferArgs() const {
return bufferArgs_;
}
private:
enum BackendType {
kUninitialized,
kSimpleIREval,
kLLVMCodeGen,
kCudaCodeGen,
kBlockCodeGen,
};
void compile();
void genInputDebugNames();
void runKernel(Stack& stack);
std::vector dimsFromSizes(const std::vector& sizes);
std::vector sizesForValue(const torch::jit::Value* v);
std::vector inferSizesForValue(const torch::jit::Value* v);
std::vector sizesFromVaryingShape(
const c10::VaryingShape& shape);
// These functions broadcast shape and also store a `hasBroadcast_` variable.
std::vector broadcastShapesMut(
const std::vector& a,
const std::vector& b);
std::vector broadcastShapesMut(
std::vector> shapes);
ExprHandle chunk(
BufPtr b,
size_t chunkIdx,
int64_t dim,
int64_t chunks,
const std::vector& axes);
ArgValue toArg(const torch::jit::Value* v) const;
ExprHandle constant(const torch::jit::Value* v);
ExprHandle tensorOrConstant(
const torch::jit::Value* v,
const std::vector& axes);
Tensor computeValue(const torch::jit::Value* v);
void bindConstant(const torch::jit::Value* v);
StmtPtr transformLoops(BackendType backendType, StmtPtr st);
std::string getCodeGenName(BackendType backendType);
std::vector prepareRunArgs(
const at::ArrayRef& inputs,
std::vector& outputs);
BackendType inferBackendTypeFromDevice(at::Device device);
Tensor bindInput(const torch::jit::Value* input);
Tensor convertOutputToCorrectStrides(torch::jit::Value* v);
// Captures the information for reduction operation nodes.
struct ReductionInfo {
std::vector reductionDims;
std::vector outputDims;
std::vector axes;
bool keepdim;
c10::optional dtype;
};
NNCLoweringFunction getCustomLoweringFor(c10::Symbol op) const;
std::unordered_map getCustomLowerings()
const {
return custom_lowerings_;
}
// Allocate memory for intermediate buffers at compile time.
// Specifically, we pre-allocate memory for intermediate buffers with static
// size and manage these buffers in the way we manage JIT constant tensors:
// push the buf args into the stack so NNC IR can access them at runtime.
void preAllocIntermediateBufs(std::unordered_set& interm_bufs);
private:
struct UnpackedTensorOptions {
c10::optional dtype;
c10::optional layout;
c10::optional device;
c10::optional pinned_memory;
UnpackedTensorOptions(const c10::TensorOptions& opts)
: dtype(optTypeMetaToScalarType(opts.dtype_opt())),
layout(opts.layout_opt()),
device(opts.device_opt()),
pinned_memory(opts.pinned_memory_opt()) {}
};
int64_t nInputs_ = 0;
std::vector bufferArgs_;
std::vector> tensorOutputSizes_;
std::vector> tensorOutputStrides_;
std::vector tensorOutputTensorOptions_;
std::unordered_set bufOutputs_;
std::unordered_map bufs_;
std::unordered_map scalars_;
std::unordered_map input_name_map_;
std::unique_ptr codegen_;
at::Device device_ = at::kCPU;
std::shared_ptr graph_;
Code code_;
bool allow_fallback_{false};
bool use_fallback_{false};
bool hasRandom_{false};
bool hasBroadcast_{false};
std::unordered_map>
known_sizes_;
std::vector unpacked_constant_tensors_;
std::vector constants_;
std::unordered_map custom_lowerings_;
bool pre_alloc_{false};
};
TORCH_API int& getTECudaPointwiseLoopLevels();
TORCH_API int& getTECudaPointwiseBlockCount();
TORCH_API int& getTECudaPointwiseBlockSize();
TORCH_API bool& getTEGenerateBlockCode();
TORCH_API bool& getTEMustUseLLVMOnCPU();
TORCH_API bool fallbackAllowed();
TORCH_API bool setFallbackAllowed(bool value);
TORCH_API bool& getCatWoConditionals();
TORCH_API bool& getOptConditionals();
TORCH_API c10::optional pickDeviceType(
const at::ArrayRef& inputs);
} // namespace tensorexpr
} // namespace jit
} // namespace torch