[TensorExpr] Add LoopNest::simplify. (#50850)

Summary: Pull Request resolved: #50850 Test Plan: Imported from OSS Reviewed By: Chillee Differential Revision: D25985085 Pulled By: ZolotukhinM fbshipit-source-id: e51709423c2c12b37b449a9d7bb22be04cda7ef1
pytorch · Jan 22, 2021 · 5f07b53 · 5f07b53
1 parent 2ba2ab9
commit 5f07b53
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp
@@ -1653,6 +1653,11 @@ bool LoopNest::hasLoopBodyFor(Tensor* t) const {
   return getLoopBodyFor(t) != nullptr;
 }
 
+Stmt* LoopNest::simplify() {
+  root_stmt_ = IRSimplifier::simplify(root_stmt_);
+  return root_stmt_;
+}
+
 Stmt* FlattenIndexes(Stmt* s) {
   IndexFlattener idx_flattener;
   return idx_flattener.flatten(s);

diff --git a/torch/csrc/jit/tensorexpr/loopnest.h b/torch/csrc/jit/tensorexpr/loopnest.h
@@ -50,6 +50,7 @@ class TORCH_API LoopNest {
   bool hasLoopBodyFor(Tensor*) const;
 
   static void vectorize(For*);
+  Stmt* simplify();
 
   bool computeInline(Stmt* s);
   bool computeInline(const Buf* b);

diff --git a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp
@@ -410,10 +410,23 @@ void initTensorExprBindings(PyObject* module) {
             self.rfactor(r, v.node(), &ins_point);
           },
           py::return_value_policy::reference)
+      .def(
+          "flatten",
+          [](const tensorexpr::LoopNest& self,
+             const std::vector<tensorexpr::For*>& loops) {
+            tensorexpr::For* flattened;
+            self.flatten(loops, &flattened);
+            return flattened;
+          },
+          py::return_value_policy::reference)
       .def(
           "reorder",
           &tensorexpr::LoopNest::reorderAxis,
           py::return_value_policy::reference)
+      .def(
+          "simplify",
+          &tensorexpr::LoopNest::simplify,
+          py::return_value_policy::reference)
       .def(
           "__str__",
           [](const tensorexpr::LoopNest& self) {