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/analysis.h (5888B)

#pragma once

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

namespace torch {
namespace jit {
namespace tensorexpr {
class HasRand : public IRVisitor {
 public:
  HasRand(StmtPtr stmt) : stmt_(stmt) {
    stmt_->accept(this);
  }

bool has_rand() const {
    return has_rand_;
  }

private:
  void visit(IntrinsicsPtr v) override {
    if (v->op_type() == IntrinsicsOp::kRand) {
      has_rand_ = true;
    } else {
      IRVisitor::visit(v);
    }
  }
  StmtPtr stmt_;
  bool has_rand_ = false;
};

template <typename Op>
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
class NodeFinder : public IRVisitor {
 public:
  void visit(NodePtr<Op> v) override {
    nodes.push_back((NodePtr<Op>)v);
    IRVisitor::visit(v);
  }

static std::vector<NodePtr<Op>> find(StmtPtr s) {
    NodeFinder<Op> nf;
    s->accept(&nf);
    return nf.nodes;
  }

static std::vector<NodePtr<Op>> find(ExprPtr e) {
    NodeFinder<Op> nf;
    e->accept(&nf);
    return nf.nodes;
  }

std::vector<NodePtr<Op>> nodes;
};

// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
class VarFinder : public IRVisitor {
 public:
  void visit(VarPtr v) override {
    vars_.insert(v);
    IRVisitor::visit(v);
  }

static std::unordered_set<VarPtr> find(StmtPtr s) {
    VarFinder nf;
    s->accept(&nf);
    return nf.vars();
  }

static std::unordered_set<VarPtr> find(ExprPtr e) {
    VarFinder nf;
    e->accept(&nf);
    return nf.vars();
  }

const std::unordered_set<VarPtr>& vars() {
    return vars_;
  }

private:
  std::unordered_set<VarPtr> vars_;
};

class BufFinder : public IRVisitor {
 public:
  void visit(BufPtr v) override {
    bufs_.insert(v);
    IRVisitor::visit(v);
  }

static std::unordered_set<BufPtr> find(StmtPtr s) {
    BufFinder nf;
    s->accept(&nf);
    return nf.bufs();
  }

static std::unordered_set<BufPtr> find(ExprPtr e) {
    BufFinder nf;
    e->accept(&nf);
    return nf.bufs();
  }

const std::unordered_set<BufPtr>& bufs() {
    return bufs_;
  }

private:
  std::unordered_set<BufPtr> bufs_;
};

// Finds all kinds of write operations to the provided Buf.
class WritesToBuf : public IRVisitor {
 public:
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  WritesToBuf(BufPtr target) : target_(target) {}

std::vector<StmtPtr> writes() {
    return writes_;
  }

static std::vector<StmtPtr> find(StmtPtr s, BufPtr b) {
    WritesToBuf finder(b);
    s->accept(&finder);
    return finder.writes();
  }

private:
  void visit(StorePtr v) override {
    if (v->buf() == target_) {
      writes_.push_back(v);
    }
  }

void visit(AtomicAddPtr v) override {
    if (v->buf() == target_) {
      writes_.push_back(v);
    }
  }

BufPtr target_;
  std::vector<StmtPtr> writes_;
};

class StmtsReadingBuf : public IRVisitor {
 public:
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  StmtsReadingBuf(BufPtr target) : target_(target) {}

std::vector<StmtPtr> reads() {
    return reads_;
  }

static std::vector<StmtPtr> find(StmtPtr s, BufPtr b) {
    StmtsReadingBuf finder(b);
    s->accept(&finder);
    return finder.reads();
  }

private:
  bool readsBuffer(StmtPtr s) {
    auto loads = NodeFinder<Load>::find(s);
    for (auto l : loads) {
      if (l->buf() == target_) {
        return true;
      }
    }
    return false;
  }

void visit(StorePtr v) override {
    if (readsBuffer(v)) {
      reads_.push_back(v);
    }
  }

void visit(LetPtr v) override {
    if (readsBuffer(v)) {
      reads_.push_back(v);
    }
  }

void visit(CondPtr v) override {
    if (readsBuffer(v)) {
      reads_.push_back(v);
    }
  }

void visit(AtomicAddPtr v) override {
    if (readsBuffer(v)) {
      reads_.push_back(v);
    }
  }

BufPtr target_;
  std::vector<StmtPtr> reads_;
};

// Traverses the IR to determine if a particular Var is modified within it.
class ModifiesVarChecker : public IRVisitor {
 public:
  ModifiesVarChecker(VarPtr v) : var_(v) {}

static bool check(StmtPtr s, VarPtr v) {
    ModifiesVarChecker checker(v);
    s->accept(&checker);
    return checker.found();
  }

bool found() {
    return found_;
  }

private:
  void visit(StorePtr v) override {
    if (v->buf()->base_handle() == var_) {
      found_ = true;
      return;
    }
    IRVisitor::visit(v);
  }

void visit(AtomicAddPtr v) override {
    if (v->buf()->base_handle() == var_) {
      found_ = true;
      return;
    }
    IRVisitor::visit(v);
  }

void visit(LetPtr v) override {
    if (v->var() == var_) {
      found_ = true;
      return;
    }
    IRVisitor::visit(v);
  }

void visit(ForPtr v) override {
    if (v->var() == var_) {
      found_ = true;
      return;
    }
    IRVisitor::visit(v);
  }

VarPtr var_;
  bool found_{false};
};

// A class that analyzes the given program relevant for Block backend
// It creates a map of multi dim buffers and their flat verions
class CreateBufferMap : public IRVisitor {
 public:
  const std::unordered_map<std::string, BufPtr>& getBufferMap() const {
    return map_input_to_tensor_bufs_;
  }

private:
  void visit(StorePtr v) override {
    auto load_node = to<Load>(v->value());
    if (load_node) {
      auto t_buf = load_node->buf();
      map_input_to_tensor_bufs_.emplace(t_buf->name_hint(), v->buf());
    } else {
      auto add_node = to<Add>(v->value());
      auto mul_node = to<Mul>(v->value());
      // This means for now, v->value() can be Add or Mul
      TORCH_INTERNAL_ASSERT(add_node || mul_node, buildErrorMessage());
      map_input_to_tensor_bufs_.emplace(v->buf()->name_hint(), v->buf());
    }
    v->value()->accept(this);
  }
  std::unordered_map<std::string, BufPtr> map_input_to_tensor_bufs_;
};

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

cmd: