FM / tensorexpr

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/jit/tensorexpr/

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/jit/tensorexpr

Name	Size	Mode	Actions
operators/	-	0755	rm
analysis.h	5888	0644	edit dl rm
block_codegen.h	4211	0644	edit dl rm
bounds_inference.h	2230	0644	edit dl rm
bounds_overlap.h	3329	0644	edit dl rm
codegen.h	6402	0644	edit dl rm
cpp_codegen.h	2278	0644	edit dl rm
cpp_intrinsics.h	719	0644	edit dl rm
cuda_codegen.h	7782	0644	edit dl rm
cuda_random.h	2642	0644	edit dl rm
dim_arg.h	884	0644	edit dl rm
eval.h	9639	0644	edit dl rm
exceptions.h	3253	0644	edit dl rm
expr.h	11588	0644	edit dl rm
external_functions.h	1274	0644	edit dl rm
external_functions_registry.h	2343	0644	edit dl rm
fwd_decls.h	2806	0644	edit dl rm
graph_opt.h	2553	0644	edit dl rm
half_support.h	5038	0644	edit dl rm
hash_provider.h	7930	0644	edit dl rm
intrinsic_symbols.h	420	0644	edit dl rm
ir.h	22622	0644	edit dl rm
ir_cloner.h	2069	0644	edit dl rm
ir_mutator.h	2010	0644	edit dl rm
ir_printer.h	3693	0644	edit dl rm
ir_simplifier.h	15090	0644	edit dl rm
ir_verifier.h	1240	0644	edit dl rm
ir_visitor.h	1825	0644	edit dl rm
kernel.h	9210	0644	edit dl rm
llvm_codegen.h	3180	0644	edit dl rm
llvm_jit.h	1965	0644	edit dl rm
loopnest.h	21599	0644	edit dl rm
mem_dependency_checker.h	13003	0644	edit dl rm
reduction.h	6742	0644	edit dl rm
registerizer.h	12498	0644	edit dl rm
stmt.h	21138	0644	edit dl rm
tensor.h	7640	0644	edit dl rm
tensorexpr_init.h	268	0644	edit dl rm
types.h	3880	0644	edit dl rm
unique_name_manager.h	940	0644	edit dl rm
var_substitutor.h	1753	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/jit/tensorexpr/codegen.h (6402B)

#pragma once

#include <ATen/ATen.h>
#include <torch/csrc/jit/tensorexpr/ir.h>
#include <torch/csrc/jit/tensorexpr/tensor.h>

namespace torch {
namespace jit {
namespace tensorexpr {

template <typename T>
class PaddedBuffer;

class TORCH_API CodeGen {
 public:
  class BufferArg;
  class CallArg;

template <typename... Ts>
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  CodeGen(StmtPtr stmt, Ts... ts)
      : stmt_(stmt), buffer_args_({BufferArg(ts)...}) {}

// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  CodeGen(
      StmtPtr stmt,
      std::vector<BufferArg> buffer_args,
      at::Device device = at::kCPU,
      std::string kernel_func_name = "func")
      : stmt_(stmt),
        buffer_args_(std::move(buffer_args)),
        device_(device),
        kernel_func_name_(std::move(kernel_func_name)) {}

virtual ~CodeGen() = default;

StmtPtr stmt() const {
    return stmt_;
  }

void set_stmt(StmtPtr s) {
    stmt_ = s;
  }

void apply_mutator(IRMutator* mutator) {
    stmt_ = stmt_->accept_mutator(mutator);
  }

void apply_visitor(IRVisitor* visitor) {
    stmt_->accept(visitor);
  }

std::vector<BufferArg>& buffer_args() {
    return buffer_args_;
  }

const std::vector<BufferArg>& buffer_args() const {
    return buffer_args_;
  }

at::Device device() {
    return device_;
  }

// This function returns the generated code as
  // a string.
  virtual std::string getCodeText(const std::string& attr = "") {
    return ("");
  }

// There are two ways to invoke the codegen:
  //  1) with a vector of CallArgs
  //  2) with a vector of raw 'void*' pointers
  //
  // The codegen knows types of all inputs from the buffer args, that's why
  // 'void*' pointers suffice.
  //
  // TODO: Eventually we might consider killing the CallArgs version, but
  // currently only LLVM codegen implements call_raw.
  virtual void call(const std::vector<CallArg>& args) = 0;
  virtual void call_raw(const std::vector<void*>& args) = 0;

virtual at::Tensor empty_strided(
      c10::IntArrayRef size,
      c10::IntArrayRef stride,
      c10::optional<c10::ScalarType> dtype_opt,
      c10::optional<c10::Layout> layout_opt,
      c10::optional<c10::Device> device_opt,
      c10::optional<bool> pin_memory_opt) {
    return at::empty_strided(
        size, stride, dtype_opt, layout_opt, device_opt, pin_memory_opt);
  }

const std::string& kernel_func_name() const {
    return kernel_func_name_;
  }

protected:
  static void* argToPtr(const BufferArg& bufferArg, const CallArg& callArg);

private:
  StmtPtr stmt_;
  std::vector<BufferArg> buffer_args_;
  at::Device device_ = at::kCPU;
  std::string kernel_func_name_ = "func";
};

class CodeGen::BufferArg {
 public:
  BufferArg(Tensor tensor) : buf_(tensor.buf()) {}
  BufferArg(const VarHandle& var) : var_(var.node()), isVar_(true) {}
  BufferArg(const BufHandle& buf) : buf_(buf.node()) {}

VarPtr var() const {
    return isVar_ ? var_ : buf_->base_handle();
  }

BufPtr buf() const {
    return buf_;
  }

bool isVar() const {
    return isVar_;
  }

Dtype dtype() const {
    return isVar_ ? var_->dtype() : buf_->dtype();
  }

private:
  VarPtr var_ = nullptr;
  BufPtr buf_ = nullptr;
  bool isVar_ = false;
};

class CodeGen::CallArg {
 public:
  template <typename T>
  CallArg(const PaddedBuffer<T>& buffer);

template <typename T>
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init,cppcoreguidelines-pro-type-const-cast)
  CallArg(const std::vector<T>& buffer)
      // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
      : data_(const_cast<T*>(buffer.data())) {}

// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  CallArg(void* ptr) : data_(ptr) {}

#define ARG_TYPE_CTOR(Type, Name)     \
  CallArg(Type v) {                   \
    memcpy(&data_, &v, sizeof(Type)); \
  }
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, ARG_TYPE_CTOR);
#undef ARG_TYPE_CTOR

void* data() const {
    return data_;
  }

#define ARG_PTR_DEFINE(Type, Name) \
  Type* Name##Ptr() const {        \
    return (Type*)&data_;          \
  }
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, ARG_PTR_DEFINE);
#undef ARG_PTR_DEFINE

private:
  void* data_;
};

// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
class RegisterCodeGenList {
 public:
  TORCH_API static RegisterCodeGenList& GetInstance() {
    static RegisterCodeGenList codegen_list;
    return codegen_list;
  }

using StmtFactoryMethod = std::function<std::unique_ptr<CodeGen>(
      StmtPtr stmt,
      const std::vector<CodeGen::BufferArg>&,
      at::Device device,
      const std::string& kernel_func_name)>;

TORCH_API StmtFactoryMethod FindStmtFactoryMethod(const std::string& name);
  RegisterCodeGenList(const RegisterCodeGenList&) = delete;
  RegisterCodeGenList& operator=(const RegisterCodeGenList&) = delete;

private:
  template <class CodeGenType>
  friend class RegisterCodeGen;
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  RegisterCodeGenList() = default;
  TORCH_API void AddStmtFactoryMethod(
      const std::string& name,
      const StmtFactoryMethod& stmt_factory_method);

std::unordered_map<std::string, StmtFactoryMethod> stmt_factory_methods_;
};

template <class CodeGenType>
class RegisterCodeGen {
 public:
  explicit RegisterCodeGen(const std::string& name) {
    RegisterCodeGenList& codegen_list = RegisterCodeGenList::GetInstance();
    codegen_list.AddStmtFactoryMethod(
        name,
        [](StmtPtr stmt,
           const std::vector<CodeGen::BufferArg>& params,
           at::Device device,
           const std::string& kernel_func_name) {
          // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
          std::unique_ptr<CodeGen> method(
              new CodeGenType(stmt, params, device, kernel_func_name));
          return method;
        });
  }
};

TORCH_API std::unique_ptr<CodeGen> CreateCodeGen(
    const std::string& name,
    StmtPtr stmt,
    const std::vector<CodeGen::BufferArg>& params,
    at::Device device = at::kCPU,
    const std::string& kernel_func_name = "func");

class TORCH_API GenericIntrinsicsExpander : public IRMutator {
 protected:
  ExprPtr mutate(IntrinsicsPtr v) override;
};

} // namespace tensorexpr
} // namespace jit
} // namespace torch

cmd: