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/var_substitutor.h (1753B)

#pragma once

#include <unordered_map>
#include <vector>

#include <torch/csrc/jit/tensorexpr/analysis.h>
#include <torch/csrc/jit/tensorexpr/ir.h>
#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
#include <torch/csrc/jit/tensorexpr/reduction.h>

namespace torch {
namespace jit {
namespace tensorexpr {

using VarMapping = std::vector<std::pair<VarPtr, ExprPtr>>;

class VarSubMutator : public IRMutator {
 public:
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
  VarSubMutator(const VarMapping& var_mapping) {
    for (auto& entry : var_mapping) {
      VarPtr key_var = entry.first;
      ExprPtr value = entry.second;
      if (!key_var) {
        throw malformed_input("missing key in VarSubMutator");
      }
      var_mapping_[key_var] = value;
    }
  }

ExprPtr mutate(VarPtr var) override {
    auto iter = var_mapping_.find(var);
    if (iter == var_mapping_.end()) {
      return var;
    }
    return iter->second;
  }

ExprPtr mutate(ReduceOpPtr var) override {
    auto body = var->body()->accept_mutator(this);
    // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
    std::vector<VarPtr> new_inner;

for (auto v : var->reduce_args()) {
      ExprPtr e = v->accept_mutator(this);
      if (VarPtr new_var = to<Var>(e)) {
        new_inner.push_back(new_var);
      } else {
        VarFinder varFinder;
        e->accept(&varFinder);
        auto varlist = varFinder.vars();
        new_inner.insert(new_inner.end(), varlist.begin(), varlist.end());
      }
    }

return alloc<ReduceOp>(body, new_inner, var->reducer());
  }

private:
  std::unordered_map<VarPtr, ExprPtr> var_mapping_;
};

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

cmd: