Lower linearized into a series of adds and muls.

taichi/transforms/simplify.cpp

@@ -661,8 +661,8 @@ class BasicBlockSimplify : public IRVisitor {
   }
 
   void visit(LinearizeStmt *stmt) override {
-    if (is_done(stmt))
-      return;
+//    if (is_done(stmt))
+//      return;
 
     if (stmt->inputs.size() && stmt->inputs.back()->is<IntegerOffsetStmt>()) {
       auto previous_offset = stmt->inputs.back()->as<IntegerOffsetStmt>();
@@ -675,23 +675,27 @@ class BasicBlockSimplify : public IRVisitor {
       offset_stmt->as<IntegerOffsetStmt>()->input = stmt;
       throw IRModified();
     }
-
-    for (int i = 0; i < current_stmt_id; i++) {
-      auto &bstmt = block->statements[i];
-      if (stmt->ret_type == bstmt->ret_type) {
-        auto &bstmt_data = *bstmt;
-        if (typeid(bstmt_data) == typeid(*stmt)) {
-          auto bstmt_ = bstmt->as<LinearizeStmt>();
-          if (identical_vectors(bstmt_->inputs, stmt->inputs) &&
-              identical_vectors(bstmt_->strides, stmt->strides)) {
-            stmt->replace_with(bstmt.get());
-            stmt->parent->erase(current_stmt_id);
-            throw IRModified();
-          }
-        }
-      }
-    }
-    set_done(stmt);
+//    set_done(stmt);
+
+    // Lower into a series of adds and muls.
+    auto sum = Stmt::make<ConstStmt>(LaneAttribute<TypedConstant>(0));
+    auto stride_product = 1;
+    for (int i = (int)stmt->inputs.size() - 1; i >= 0; i--) {
+      auto stride_stmt = Stmt::make<ConstStmt>(LaneAttribute<TypedConstant>(stride_product));
+      auto mul = Stmt::make<BinaryOpStmt>(BinaryOpType::mul, stmt->inputs[i], stride_stmt.get());
+      auto newsum = Stmt::make<BinaryOpStmt>(BinaryOpType::add, sum.get(), mul.get());
+      stmt->insert_before_me(std::move(sum));
+      sum = std::move(newsum);
+      stmt->insert_before_me(std::move(stride_stmt));
+      stmt->insert_before_me(std::move(mul));
+      stride_product *= stmt->strides[i];
+    }
+    stmt->replace_with(sum.get());
+    stmt->insert_before_me(std::move(sum));
+    stmt->parent->erase(stmt);
+    // get types of adds and muls
+    irpass::typecheck(stmt->parent);
+    throw IRModified();
   }
 
   void visit(SNodeLookupStmt *stmt) override {

tests/cpp/test_alg_simp.cpp

@@ -1,6 +1,5 @@
 #include <taichi/lang.h>
 #include <taichi/testing.h>
-#include <taichi/program.h>
 
 TLANG_NAMESPACE_BEGIN
 

tests/cpp/test_simplify.cpp

@@ -0,0 +1,40 @@
+#include <taichi/lang.h>
+#include <taichi/testing.h>
+
+TLANG_NAMESPACE_BEGIN
+
+// Basic tests within a basic block
+
+TI_TEST("simplify_linearized_with_trivial_inputs") {
+  auto block = std::make_unique<Block>();
+
+  auto get_root = block->push_back<GetRootStmt>();
+  auto linearized_empty =
+      block->push_back<LinearizeStmt>(std::vector<Stmt*>(), std::vector<int>());
+  SNode root(0, SNodeType::root);
+  root.insert_children(SNodeType::dense);
+  auto lookup =
+      block->push_back<SNodeLookupStmt>(&root, get_root, linearized_empty, false, std::vector<Stmt*>());
+  auto get_child = block->push_back<GetChStmt>(lookup, 0);
+  auto zero = block->push_back<ConstStmt>(TypedConstant(0));
+  auto linearized_zero =
+      block->push_back<LinearizeStmt>(std::vector<Stmt*>(2, zero), std::vector<int>({8, 4}));
+  auto lookup2 =
+      block->push_back<SNodeLookupStmt>(root.ch[0].get(), get_child, linearized_zero, true, std::vector<Stmt*>());
+
+  irpass::typecheck(block.get());
+  // irpass::print(block.get());
+  TI_CHECK(block->size() == 7);
+
+  irpass::simplify(block.get());  // should lower linearized
+  // irpass::print(block.get());
+  // TI_CHECK(block->size() == 11);  // not required to check size here
+
+  irpass::constant_fold(block.get());
+  irpass::alg_simp(block.get(), CompileConfig());
+  irpass::die(block.get());  // should eliminate consts
+  // irpass::print(block.get());
+  TI_CHECK(block->size() == 5);  // get root, const 0, lookup, get child, lookup
+}
+
+TLANG_NAMESPACE_END