/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/tensor.h (7640B)
#pragma once
#include
#include
#include
#include
#include
#include
namespace torch {
namespace jit {
namespace tensorexpr {
class TORCH_API Tensor {
public:
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
Tensor(BufPtr buf, const std::vector& args, ExprPtr body)
: buf_(buf) {
stmt_ = constructStmt(args, body, {}, {});
}
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
Tensor(
BufPtr buf,
const std::vector& args,
const std::vector& reduce_dims,
const std::vector& reduce_args,
ExprPtr body)
: buf_(buf) {
stmt_ = constructStmt(args, body, reduce_dims, reduce_args);
}
Tensor(BufPtr buf, StmtPtr stmt) : buf_(buf), stmt_(stmt) {}
BufPtr buf() const {
return buf_;
}
StmtPtr stmt() const {
return stmt_;
}
template
inline ExprHandle load(const std::vector& args) const;
template
inline ExprHandle load(const Ts&... ts) const;
private:
StmtPtr constructStmt(
const std::vector& args,
ExprPtr body,
const std::vector& reduce_dims,
const std::vector& reduce_args) const;
BufPtr buf_;
StmtPtr stmt_;
};
TORCH_API Tensor Compute(
const std::string& func_name,
const std::vector& dim_args,
const std::function& body_func);
TORCH_API Tensor Compute(
const std::string& func_name,
const std::vector& dim_args,
const std::function&
body_func);
TORCH_API Tensor Compute(
const std::string& func_name,
const std::vector& dim_args,
const std::function<
ExprHandle(const VarHandle&, const VarHandle&, const VarHandle&)>&
body_func);
TORCH_API Tensor Compute(
const std::string& func_name,
const std::vector& dim_args,
const std::function& body_func);
TORCH_API Tensor Compute(
const std::string& func_name,
const std::vector& dim_args,
const std::function&)>& body_func);
inline void unpack_dim_args(
const std::vector& dim_args,
std::vector* dims,
std::vector* vars) {
dims->clear();
vars->clear();
for (const DimArg& dim_arg : dim_args) {
ExprPtr expr = dim_arg.dim().node();
dims->push_back(expr);
vars->push_back(alloc(
dim_arg.name_hint(),
expr->dtype().scalar_type() == ScalarType::Long ? kLong : kInt));
}
}
// Handle reductions over a Reducer and a body_func which produces values.
template
Tensor Reduce(
const std::string& func_name,
const std::vector& dim_args,
const Reducer& reducer,
const InitFunc& init_func,
const BodyFunc& body_func,
const std::vector& reduce_args) {
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector dims;
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector vars;
unpack_dim_args(dim_args, &dims, &vars);
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector reduce_dims;
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector reduce_vars;
unpack_dim_args(reduce_args, &reduce_dims, &reduce_vars);
// If reduce_vars is empty, then it's not a reduction, but rather a simple
// copy
if (reduce_vars.empty()) {
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
ExprPtr body =
Reducer::getReduceBody(body_func, VarVectorToVarHandleVector(vars))
.node();
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
BufPtr func_result = alloc(func_name, dims, body->dtype());
return Tensor(func_result, vars, body);
}
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector all_vars;
all_vars.insert(all_vars.end(), vars.begin(), vars.end());
all_vars.insert(all_vars.end(), reduce_vars.begin(), reduce_vars.end());
ExprHandle body =
Reducer::getReduceBody(body_func, VarVectorToVarHandleVector(all_vars));
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector output_args(vars.begin(), vars.end());
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
ExprPtr init_expr = alloc(
body.dtype(), init_func(VarVectorToVarHandleVector(vars)).node());
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
BufPtr func_result = alloc(func_name, dims, body.dtype(), init_expr);
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
ReduceOpPtr reduce_op = reducer(func_result, body, output_args, reduce_vars);
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
Tensor t = Tensor(func_result, vars, reduce_dims, reduce_vars, reduce_op);
return t;
}
template
Tensor Reduce(
const std::string& func_name,
const std::vector& dim_args,
const Reducer& reducer,
const BodyFunc& body_func,
const std::vector& reduce_args) {
return Reduce(
func_name,
dim_args,
reducer,
[&](ParameterList p) { return ExprHandle(reducer.initializer()); },
body_func,
reduce_args);
}
// Overload which allows inline lambda functions for the body_func.
template
Tensor Reduce(
const std::string& func_name,
const std::vector& dim_args,
const Reducer& reducer,
const BodyFunc&& body_func,
const std::vector& reduce_args) {
return Reduce(func_name, dim_args, reducer, body_func, reduce_args);
}
TORCH_API Tensor Reduce(
const std::string& name,
const std::vector& dim_args,
const Reducer& reducer,
const BufHandle& buffer,
const std::vector& reduce_args);
// Overload for the common case of all dimensions of a prevously Computed
// Tensor.
TORCH_API Tensor Reduce(
const std::string& func_name,
const std::vector& dim_args,
const Reducer& reducer,
Tensor tensor,
const std::vector& reduce_args);
template
inline ExprHandle Tensor::load(const Ts&... ts) const {
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector params({ExprHandle(ts)...});
return Load::make(BufHandle(this->buf()), params);
}
template
inline ExprHandle Tensor::load(const std::vector& args) const {
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector params(args.begin(), args.end());
return Load::make(BufHandle(this->buf()), params);
}
template
inline ExprHandle BufHandle::load(const Ts&... ts) const {
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector params({ExprHandle(ts)...});
return ExprHandle(alloc(node(), ExprHandleVectorToExprVector(params)));
}
template
inline ExprHandle BufHandle::load(const std::vector& args) const {
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector params(args.begin(), args.end());
return ExprHandle(alloc(node(), ExprHandleVectorToExprVector(params)));
}
inline ExprHandle BufHandle::load(const std::vector& args) const {
return this->template load(args);
}
} // namespace tensorexpr
} // namespace jit
} // namespace torch