[Fix] Remove CastRemover for Tensor Indices (issue cornell-zhang#386) (cornell-zhang#453)

* [Test] Add test case for issue cornell-zhang#386

* [Fix] Remove CastRemover for tensor indices and make sure index is i32 type in LLVM backend (Fix cornell-zhang#386)

* [Fix] Fix issue with simplying index expression in generate_reuse_buffer

* [Fix] Add type checking in CastRemover

* [Fix] Move CastRemover to ir_util.h

Author: zzzDavid

python/heterocl/tensor.py

@@ -148,7 +148,6 @@ def __setitem__(self, indices, expr):
         if not isinstance(indices, tuple):
             indices = (indices,)
         indices = self.indices + indices
-        indices = util.CastRemover().mutate(indices)
         index, bit, _ = util.get_index(self.tensor.shape, indices, 0)
         if not Stage.get_len():
             raise TensorError("Cannot set tensor elements without compute APIs")
@@ -238,7 +237,6 @@ def shape(self):
     def asnode(self):
         if len(self.indices) < len(self.tensor.shape):
             raise TensorError("Accessing a slice of tensor is not allowed")
-        self.indices = util.CastRemover().mutate(self.indices)
         index, bit, _ = util.get_index(self.tensor.shape, self.indices, 0)
         if bit is None:
             return _make.Load(self._dtype, self.tensor.buf.data, index)
@@ -343,20 +341,17 @@ def __repr__(self):
         return "Tensor('" + self.name + "', " + str(self.shape) + ", " + str(self.dtype) + ")"
 
     def __getitem__(self, indices):
-        indices = util.CastRemover().mutate(indices)
         if Stage.get_len():
             Stage.get_current().input_stages.add(self.last_update)
         if not isinstance(indices, tuple):
             indices = (indices,)
         return TensorSlice(self, indices)
 
     def __setitem__(self, indices, expr):
-        indices = util.CastRemover().mutate(indices)
         Stage.get_current().input_stages.add(self.last_update)
         Stage.get_current().lhs_tensors.add(self)
         if not isinstance(indices, tuple):
             indices = (indices,)
-        indices = util.CastRemover().mutate(indices)
         if len(indices) < len(self.shape):
             raise TensorError("Accessing a slice of tensor is not allowed")
         else:

tests/issues/test_issue_386.py

@@ -0,0 +1,23 @@
+import heterocl as hcl
+
+def test_cast_removal():
+    hcl.init()
+    
+    A = hcl.placeholder((10,10), dtype=hcl.UInt(16), name="A")
+    B = hcl.placeholder((10,10), dtype=hcl.Int(16), name="B")
+
+    def algo(A, B):
+        def f_mutate(i,j):
+            factor = hcl.scalar(B[0][0][13:11], name="factor")
+            idx = hcl.scalar(B[0][0][11:0], dtype=hcl.UInt(16), name="idx")
+            idx += i * hcl.cast(hcl.UInt(16), factor.v) 
+            A[idx][j] = B[idx][j]
+        bound = hcl.scalar(5, dtype=hcl.Int(32))
+        domain = (hcl.cast(hcl.UInt(32), bound.v), hcl.cast(hcl.UInt(32), bound.v))
+        hcl.mutate(domain, f_mutate)
+    
+    s = hcl.create_schedule([A, B], algo)
+    f = hcl.build(s, target="vhls")
+
+if __name__ == '__main__':
+    test_cast_removal()

tvm/src/codegen/llvm/codegen_llvm.cc

@@ -972,7 +972,11 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const Load* op) {
   Type t = op->type;
   bool is_volatile = volatile_buf_.count(op->buffer_var.get());
   llvm::Value* buffer = MakeValue(op->buffer_var);
+  // Make sure index is int32 type since we removed CastRemover pass.
+  // check op->index's datatype, if it is not int32, cast it to int32
   llvm::Value* index = MakeValue(op->index);
+  if (op->index.type() != Int(32))
+    index = CreateCast(op->index.type(), Int(32), index);
 
   if (t.lanes() == 1) {
     int alignment, native_bits;
@@ -1257,6 +1261,11 @@ void CodeGenLLVM::VisitStmt_(const Store* op) {
   llvm::Value* index = MakeValue(op->index);
   llvm::Value* value = MakeValue(op->value);
 
+  // Make sure index is int32 type since we removed CastRemover pass.
+  // check op->index's datatype, if it is not int32, cast it to int32
+  if (op->index.type() != Int(32))
+    index = CreateCast(op->index.type(), Int(32), index);
+
   if (t.lanes() == 1) {
     int alignment, native_bits;
     GetAlignment(t, op->buffer_var.get(), op->index, &alignment, &native_bits);

tvm/src/pass/generate_reuse_buffer.cc

@@ -8,6 +8,7 @@
 #include <tvm/ir_pass.h>
 #include <tvm/ir_visitor.h>
 #include <tvm/operation.h>
+#include "ir_util.h"
 
 namespace TVM {
 namespace ir {
@@ -222,7 +223,8 @@ class ReuseBufferInserter final : public IRMutator {
         // check if the bounde is constant
         // e.g. x+r => diff_expr = 10
         // e.g. y+c => diff_expr = 3
-        Expr diff_expr = Simplify(max_expr - min_expr + 1);
+        CastRemover castRemover;
+        Expr diff_expr = Simplify(castRemover.Mutate(max_expr - min_expr) + 1);
         if (!is_const(diff_expr))  // e.g. y*(y+c) would be illegal
           LOG(FATAL) << "Irregular access pattern is not yet supported";
         // check if the specified axis is reused by running the next iteration
@@ -231,7 +233,7 @@ class ReuseBufferInserter final : public IRMutator {
         // first check if the axis is the specified reuse axis
         // e.g. y => y+1
         Expr next_min = substitute(next_subst, min_expr);
-        Expr next_diff = Simplify(next_min - min_expr);
+        Expr next_diff = Simplify(castRemover.Mutate(next_min - min_expr));
         if (!is_const(next_diff))  // e.g. y*y+c would be illegal
           LOG(FATAL) << "Irregular access pattern is not yet supported";
         // then check if we there is reuse in this axis
@@ -262,7 +264,10 @@ class ReuseBufferInserter final : public IRMutator {
       std::vector<VarExpr> reuse_loop_vars;
       for (size_t dim = 0; dim < ndim; dim++) {
         Expr index = min_list[dim];
+        CastRemover castRemover;
+        index = castRemover.Mutate(index);
         Expr reuse_index = Simplify(substitute(null_axis_subst_, index));
+
         // create a new variable if the shape is not one
         if (!is_one(reuse_shape[dim])) {
           // TODO(Sean): fix the name
@@ -280,7 +285,9 @@ class ReuseBufferInserter final : public IRMutator {
               }
             }
           }
+
           Expr rhs = substitute(reuse_index, new_loop_var, index);
+
           // special case when the reuse index is 0
           if (is_zero(reuse_index) && dim == static_cast<size_t>(reuse))
             rhs = rhs + new_loop_var;

tvm/src/pass/ir_util.h

@@ -8,6 +8,7 @@
 
 #include <tvm/ir.h>
 #include <tvm/runtime/device_api.h>
+#include <tvm/ir_mutator.h>
 #include <vector>
 
 namespace TVM {
@@ -154,6 +155,124 @@ inline int GetTempAllocaAlignment(Type type, int32_t const_size) {
   }
   return align;
 }
+
+// Remove cast in binary expressions
+// Example: (cast(x) + cast(y)) -> (x + y)
+// Usage: CastRemover castRemover;
+//        Expr expr = castRemover.Mutate(expr);
+class CastRemover final : public IRMutator {
+ public:
+  CastRemover() {}
+
+  Expr Mutate_(const Cast* op, const Expr& e) {
+    return op->value;
+  }
+
+  Expr Mutate_(const Add* op, const Expr& e) {
+    Expr a = this->Mutate(op->a);
+    Expr b = this->Mutate(op->b);
+    if (const Cast* ca = a.as<Cast>()) {
+      a = ca->value;
+    }
+    if (const Cast* cb = b.as<Cast>()) {
+      b = cb->value;
+    }
+    if (a.type() != b.type())
+      LOG(FATAL) << "CastRemover: type mismatch "
+        << a.type() << " vs " << b.type();
+    return Add::make(a, b);
+  }
+
+  Expr Mutate_(const Sub* op, const Expr& e) {
+    Expr a = this->Mutate(op->a);
+    Expr b = this->Mutate(op->b);
+    if (const Cast* ca = a.as<Cast>()) {
+      a = ca->value;
+    }
+    if (const Cast* cb = b.as<Cast>()) {
+      b = cb->value;
+    }
+    if (a.type() != b.type())
+      LOG(FATAL) << "CastRemover: type mismatch "
+        << a.type() << " vs " << b.type();
+    return Sub::make(a, b);
+  }
+
+  Expr Mutate_(const Mul* op, const Expr& e) {
+    Expr a = this->Mutate(op->a);
+    Expr b = this->Mutate(op->b);
+    if (const Cast* ca = a.as<Cast>()) {
+      a = ca->value;
+    }
+    if (const Cast* cb = b.as<Cast>()) {
+      b = cb->value;
+    }
+    if (a.type() != b.type())
+      LOG(FATAL) << "CastRemover: type mismatch "
+        << a.type() << " vs " << b.type();
+    return Mul::make(a, b);
+  }
+
+  Expr Mutate_(const Div* op, const Expr& e) {
+    Expr a = this->Mutate(op->a);
+    Expr b = this->Mutate(op->b);
+    if (const Cast* ca = a.as<Cast>()) {
+      a = ca->value;
+    }
+    if (const Cast* cb = b.as<Cast>()) {
+      b = cb->value;
+    }
+    if (a.type() != b.type())
+      LOG(FATAL) << "CastRemover: type mismatch "
+        << a.type() << " vs " << b.type();
+    return Div::make(a, b);
+  }
+
+  Expr Mutate_(const Mod* op, const Expr& e) {
+    Expr a = this->Mutate(op->a);
+    Expr b = this->Mutate(op->b);
+    if (const Cast* ca = a.as<Cast>()) {
+      a = ca->value;
+    }
+    if (const Cast* cb = b.as<Cast>()) {
+      b = cb->value;
+    }
+    if (a.type() != b.type())
+      LOG(FATAL) << "CastRemover: type mismatch "
+        << a.type() << " vs " << b.type();
+    return Mod::make(a, b);
+  }
+
+  Expr Mutate_(const Min* op, const Expr& e) {
+    Expr a = this->Mutate(op->a);
+    Expr b = this->Mutate(op->b);
+    if (const Cast* ca = a.as<Cast>()) {
+      a = ca->value;
+    }
+    if (const Cast* cb = b.as<Cast>()) {
+      b = cb->value;
+    }
+    if (a.type() != b.type())
+      LOG(FATAL) << "CastRemover: type mismatch "
+        << a.type() << " vs " << b.type();
+    return Min::make(a, b);
+  }
+
+  Expr Mutate_(const Max* op, const Expr& e) {
+    Expr a = this->Mutate(op->a);
+    Expr b = this->Mutate(op->b);
+    if (const Cast* ca = a.as<Cast>()) {
+      a = ca->value;
+    }
+    if (const Cast* cb = b.as<Cast>()) {
+      b = cb->value;
+    }
+    if (a.type() != b.type())
+      LOG(FATAL) << "CastRemover: type mismatch "
+        << a.type() << " vs " << b.type();
+    return Max::make(a, b);
+  }
+};
 }  // namespace ir
 }  // namespace TVM
 #endif  // PASS_IR_UTIL_H_