You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/03/18 14:36:18 UTC
[GitHub] [incubator-tvm] hzfan opened a new pull request #5092: [PASS] dtype
rewrite for indexing variables
hzfan opened a new pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092
Changes:
- enable indexing with i64 vars, so that large tensors with more than 2^32 elements can be properly indexed.
- narrow i64 index which trivially fits into i32 to i32.
Some background:
https://discuss.tvm.ai/t/rfc-support-for-large-tensors/5643
Take the following as an example:
```
A = te.placeholder((m, n), name='A')
B = te.placeholder((m, n), name='B')
C = te.compute((m, n), lambda *idx: A[idx] + B[idx])
```
`m, n = te.var(’m’, dtype=‘int64’), te.var(’n’, dtype=‘int64’)` yields
```
produce compute {
for (i0.int64, (int64)0, m.int64) {
for (i1.int64, (int64)0, n.int64) {
compute[((i0.int64*stride.int64) + (i1.int64*stride.int64))] = (A[((i0.int64*stride.int64) + (i1.int64*stride.int64))] + B[((i0.int64*stride.int64) + (i1.int64*stride.int64))])
}
}
}
```
`m, n = tvm.tir.const(2, dtype="int64"), tvm.tir.const(2, dtype="int64")` yields
```
produce compute {
for (i0.int32, 0, 2) {
for (i1.int32, 0, 2) {
compute[((i0.int32*2) + i1.int32)] = (A[((i0.int32*2) + i1.int32)] + B[((i0.int32*2) + i1.int32)])
}
}
}
```
@yzhliu Could you review?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395444098
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,326 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
+ bits_ = bits > bits_ ? bits : bits_;
+ StmtExprVisitor::VisitExpr(e);
+ bits_ = tmp;
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vextent_[op->loop_var.as<VarNode>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vextent_[iv->var.as<VarNode>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vextent_.find(op) != vextent_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const IntImmNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ // the narrowed datatype of Var and IntImm
+ std::unordered_map<const PrimExprNode*, DataType> vmap;
+
+ protected:
+ // internal analyzer
+ arith::Analyzer analyzer_;
+
+ private:
+ // the maximum possible bit of the current expression's return dtype
+ int bits_;
+ // the extent of vars to be rewritten
+ std::unordered_map<const VarNode*, DataType> vextent_;
+};
+
+class DataTypeRewriter : public StmtExprMutator {
+ public:
+ Stmt operator()(Stmt s) {
+ visitor_(s);
+ for (auto i = visitor_.vmap.begin(), last = visitor_.vmap.end(); i != last;) {
+ PrimExpr e = GetRef<PrimExpr>(i->first);
+ if (e.dtype() == i->second) {
+ i = visitor_.vmap.erase(i);
+ } else {
+ ++i;
+ }
+ }
+ return VisitStmt(s);
+ }
+
+ Stmt VisitStmt_(const StoreNode* op) final {
+ PrimExpr value = this->VisitExpr(op->value);
+ is_index_ = true;
+ PrimExpr index = this->VisitExpr(op->index);
+ is_index_ = false;
+ Stmt s = StoreNode::make(op->buffer_var,
+ op->value,
+ index,
+ op->predicate);
+ return StmtExprMutator::VisitStmt_(s.as<StoreNode>());
+ }
+
+ Stmt VisitStmt_(const ForNode* op) final {
+ Stmt s = StmtExprMutator::VisitStmt_(op);
+ op = s.as<ForNode>();
+ PrimExpr e = VisitExpr(op->loop_var);
+ Var var = Downcast<Var, PrimExpr>(e);
+ return ForNode::make(var, cast(var.dtype(), op->min), cast(var.dtype(), op->extent),
+ op->for_type, op->device_api, op->body);
+ }
+
+ Stmt VisitStmt_(const AttrStmtNode* op) final {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ Stmt s = StmtExprMutator::VisitStmt_(op);
+ op = s.as<AttrStmtNode>();
+ const IterVarNode* iv = op->node.as<IterVarNode>();
+ PrimExpr e = VisitExpr(iv->var);
+ Var var = Downcast<Var, PrimExpr>(e);
+ if (ivmap_.find(iv) == ivmap_.end()) {
+ ivmap_[iv] = IterVarNode::make(iv->dom, var, iv->iter_type, iv->thread_tag);
+ }
+ return AttrStmtNode::make(
+ ivmap_[iv],
+ op->attr_key,
+ cast(var.dtype(), op->value),
+ op->body);
+ }
+ return StmtExprMutator::VisitStmt_(op);
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) final {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ if (vmap_.find(op) == vmap_.end()) {
+ vmap_[op] = Var(op->name_hint, visitor_.vmap[op]);
+ }
+ return vmap_[op];
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const SizeVarNode* op) final {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ if (vmap_.find(op) == vmap_.end()) {
+ vmap_[op] = SizeVar(op->name_hint, visitor_.vmap[op]);
+ }
+ return vmap_[op];
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) final {
+ is_index_ = true;
+ PrimExpr index = this->VisitExpr(op->index);
+ is_index_ = false;
+ PrimExpr e = LoadNode::make(op->dtype, op->buffer_var, index, op->predicate);
+ return StmtExprMutator::VisitExpr_(e.as<LoadNode>());
+ }
+
+ PrimExpr VisitExpr_(const IntImmNode* op) final {
+ if (is_index_) {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ return IntImm(visitor_.vmap[op], op->value);
+ }
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const CastNode* op) final {
+ if (is_index_ && visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ PrimExpr e = StmtExprMutator::VisitExpr_(op);
+ const CastNode* new_op = e.as<CastNode>();
+ return CastNode::make(visitor_.vmap[op], new_op->value);
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) final;
+ PrimExpr VisitExpr_(const SubNode* op) final;
+ PrimExpr VisitExpr_(const MulNode* op) final;
+ PrimExpr VisitExpr_(const DivNode* op) final;
+ PrimExpr VisitExpr_(const ModNode* op) final;
+ PrimExpr VisitExpr_(const FloorDivNode* op) final;
+ PrimExpr VisitExpr_(const FloorModNode* op) final;
+ PrimExpr VisitExpr_(const MinNode* op) final;
+ PrimExpr VisitExpr_(const MaxNode* op) final;
+ PrimExpr VisitExpr_(const EQNode* op) final;
+ PrimExpr VisitExpr_(const NENode* op) final;
+ PrimExpr VisitExpr_(const LTNode* op) final;
+ PrimExpr VisitExpr_(const LENode* op) final;
+ PrimExpr VisitExpr_(const GTNode* op) final;
+ PrimExpr VisitExpr_(const GENode* op) final;
+ PrimExpr VisitExpr_(const CallNode* op) final;
+
+ private:
+ // the internal visitor to deduce the narrowed dtype
+ DataTypeVisitor visitor_;
+ // a map from Var before rewrite to that after rewrite,
+ // ensures one old Var maps to exactly one new Var
+ std::unordered_map<const VarNode*, Var> vmap_;
+ // a map from IterVar before rewrite to that after rewrite,
+ // ensures one old IterVar maps to exactly one new IterVar
+ std::unordered_map<const IterVarNode*, IterVar> ivmap_;
+ bool is_index_{false};
+};
+
+#define DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC) \
+ PrimExpr DataTypeRewriter::VisitExpr_(const OP* op) { \
+ PrimExpr a = this->VisitExpr(op->a); \
+ PrimExpr b = this->VisitExpr(op->b); \
+ if (a.same_as(op->a) && \
+ b.same_as(op->b)) { \
+ return GetRef<PrimExpr>(op); \
+ } else { \
+ return FUNC(a, b); \
+ } \
+ }
+
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(AddNode, operator+)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(SubNode, operator-)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MulNode, operator*)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(DivNode, div)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(ModNode, truncmod)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(FloorDivNode, floordiv)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(FloorModNode, floormod)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MinNode, min)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MaxNode, max)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(EQNode, operator==)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(NENode, operator!=)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LTNode, operator <)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LENode, operator<=)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GTNode, operator >)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GENode, operator>=)
+
+PrimExpr DataTypeRewriter::VisitExpr_(const CallNode* op) {
+ PrimExpr e = StmtExprMutator::VisitExpr_(op);
+ op = e.as<CallNode>();
+ if (op->call_type == CallNode::PureIntrinsic) {
+ if (op->name == intrinsic::tvm_if_then_else) {
+ return if_then_else(op->args[0], op->args[1], op->args[2]);
Review comment:
what is the difference between this and simply return `e`?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r401299065
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
Review comment:
memo can have unintended consequences if the vars can be bound to different context dependent info(e.g. `if (x<10) {x+1; } else x;` `x<10` is only effective in the then branch.
I would say perhaps we could have another API to pass in a unordered map, and ask the analyzer to record every intermediate steps into the map
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395156871
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
Review comment:
Added a term `\brief`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397930633
##########
File path: tests/python/unittest/test_tir_pass_narrow_datatype.py
##########
@@ -0,0 +1,128 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+from tvm import te
+
+
+def lower(sch, args):
+ binds = {}
+ arg_list = []
+ for x in args:
+ if isinstance(x, te.tensor.Tensor):
+ buf = tvm.tir.decl_buffer(x.shape, dtype=x.dtype, name=x.name)
+ assert x not in binds
+ binds[x] = buf
+ arg_list.append(buf)
+ else:
+ raise ValueError("args must be Tensor, Buffer or Var")
+ bounds = te.schedule.InferBound(sch)
+ stmt = te.schedule.ScheduleOps(sch, bounds)
+ stmt = tvm.tir.ir_pass.StorageFlatten(stmt, binds, 64, False)
+ stmt = tvm.tir.ir_pass.NarrowDataType(stmt, 32)
+ return stmt
+
+
+def test_const():
+ def test(m, n, dtype):
Review comment:
test -> check, because test can sometimes be picked up by test tools like nose
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395833364
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,363 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
Review comment:
Consider use std::swap
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395451760
##########
File path: include/tvm/tir/ir_pass.h
##########
@@ -385,6 +385,13 @@ Stmt DecorateDeviceScope(Stmt stmt);
*/
Stmt HoistIfThenElse(Stmt stmt);
+/*!
+ * \brief Narrow down PrimExpr datatype in stmt
+ * \param stmt The stmt to do datatype rewrite
+ * \return Transformed stmt.
+ */
+Stmt DataTypeRewrite(Stmt stmt);
Review comment:
is it better to name it "DataTypeNarrow" ?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395721838
##########
File path: src/tir/pass/unroll_loop.cc
##########
@@ -160,7 +160,9 @@ class LoopUnroller : public StmtExprMutator {
PrimExpr extent = tir::Simplify(op->extent);
const IntImmNode *v1 = extent.as<IntImmNode>();
int value = -1;
- if (v1 != nullptr) {
+ // integers that do not fit in int32_t are treated as symbolic,
+ // as it's impossible to unroll such large loops
Review comment:
I'm not sure. A warning here may be raised for loops that are not marked as `unroll`. So every loop with a i64 extent raises the warning here, even when the user does not intend to unroll it.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r400253959
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
+ (bound->max_value <= ubound && bound->min_value >= lbound)) {
+ bits = target_bits_;
+ }
+ int tmp = bits > bits_ ? bits : bits_;
+ std::swap(bits_, tmp);
+ StmtExprVisitor::VisitExpr(e);
+ std::swap(bits_, tmp);
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vextent_[op->loop_var.as<VarNode>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vextent_[iv->var.as<VarNode>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vextent_.find(op) != vextent_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
Review comment:
Thanks for the example. Comments added.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r401289944
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
Review comment:
@hzfan I think it is good. but why not always doing memorization? @tqchen
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen edited a comment on issue #5092: [TIR][PASS]
dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen edited a comment on issue #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#issuecomment-603902373
Some more comments, mainly wrt to clarity of the code, test coverage and efficiency concerns. Thanks for bringing in the PR. given that this is critical to a lot of the codebase, let us try to https://docs.tvm.ai/contribute/code_review.html#hold-the-highest-standard :) Let us work to polish it to the best state. Thanks @hzfan for good work so far
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395731195
##########
File path: include/tvm/tir/ir_pass.h
##########
@@ -385,6 +385,13 @@ Stmt DecorateDeviceScope(Stmt stmt);
*/
Stmt HoistIfThenElse(Stmt stmt);
+/*!
+ * \brief Narrow down PrimExpr datatype in stmt
+ * \param stmt The stmt to do datatype rewrite
+ * \return Transformed stmt.
+ */
+Stmt DataTypeRewrite(Stmt stmt);
Review comment:
Renamed.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Hzfengsy commented on a change in pull request
#5092: [TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r402422599
##########
File path: src/tir/pass/unroll_loop.cc
##########
@@ -160,7 +160,9 @@ class LoopUnroller : public StmtExprMutator {
PrimExpr extent = tir::Simplify(op->extent);
const IntImmNode *v1 = extent.as<IntImmNode>();
int value = -1;
- if (v1 != nullptr) {
+ // integers that do not fit in int32_t are treated as symbolic,
+ // as it's impossible to unroll such large loops
+ if (v1 != nullptr && v1->value <= std::numeric_limits<int>::max()) {
Review comment:
That's fine
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r394812364
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
+ bits_ = bits > bits_ ? bits : bits_;
+ StmtExprVisitor::VisitExpr(e);
+ bits_ = tmp;
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vset_.insert(op->loop_var.as<Object>());
+ vextent_[op->loop_var.as<Object>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vset_.find(op) != vset_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const IntImmNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ // the narrowed datatype of Var and IntImm
+ std::unordered_map<const Object*, DataType> vmap;
+
+ protected:
+ // internal analyzer
+ arith::Analyzer analyzer_;
+
+ private:
+ // the maximum bits of all containing expressions
Review comment:
to be more precise, this is the maximum bits of the current expression return dtype, right?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395833195
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,363 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
Review comment:
Allow us to parameterize the target type(instead of hardcode 32 bits)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on issue #5092: [TIR][PASS] dtype
rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on issue #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#issuecomment-607938301
Thanks @hzfan for keep improving the code and maintain a high standard. Thanks @yzhliu @Hzfengsy for helpful reviews, this PR is now merged
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r394794367
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
Review comment:
briefly describe what does this file do.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r394803334
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
+ bits_ = bits > bits_ ? bits : bits_;
+ StmtExprVisitor::VisitExpr(e);
+ bits_ = tmp;
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vset_.insert(op->loop_var.as<Object>());
+ vextent_[op->loop_var.as<Object>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vset_.find(op) != vset_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const IntImmNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ // the narrowed datatype of Var and IntImm
+ std::unordered_map<const Object*, DataType> vmap;
+
+ protected:
+ // internal analyzer
+ arith::Analyzer analyzer_;
+
+ private:
+ // the maximum bits of all containing expressions
+ int bits_;
+ // the vars to be rewritten
+ std::unordered_set<const Object*> vset_;
Review comment:
is it necessary? can we simply check `vextent_`?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397929962
##########
File path: tests/python/unittest/test_tir_pass_narrow_datatype.py
##########
@@ -0,0 +1,128 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+from tvm import te
+
+
+def lower(sch, args):
+ binds = {}
+ arg_list = []
+ for x in args:
+ if isinstance(x, te.tensor.Tensor):
+ buf = tvm.tir.decl_buffer(x.shape, dtype=x.dtype, name=x.name)
+ assert x not in binds
+ binds[x] = buf
+ arg_list.append(buf)
+ else:
+ raise ValueError("args must be Tensor, Buffer or Var")
+ bounds = te.schedule.InferBound(sch)
+ stmt = te.schedule.ScheduleOps(sch, bounds)
+ stmt = tvm.tir.ir_pass.StorageFlatten(stmt, binds, 64, False)
+ stmt = tvm.tir.ir_pass.NarrowDataType(stmt, 32)
Review comment:
It would be great if we can add more test cases:
- Please also consider use ir_builde to directly build loops
- Have testcase that narrows to i16
- Have a loop variable occurs in multiple expressions, one expr overflows, another does not
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395155241
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
+ bits_ = bits > bits_ ? bits : bits_;
+ StmtExprVisitor::VisitExpr(e);
+ bits_ = tmp;
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vset_.insert(op->loop_var.as<Object>());
+ vextent_[op->loop_var.as<Object>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vset_.find(op) != vset_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const IntImmNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ // the narrowed datatype of Var and IntImm
+ std::unordered_map<const Object*, DataType> vmap;
+
+ protected:
+ // internal analyzer
+ arith::Analyzer analyzer_;
+
+ private:
+ // the maximum bits of all containing expressions
+ int bits_;
+ // the vars to be rewritten
+ std::unordered_set<const Object*> vset_;
Review comment:
Agree. `vset_` removed
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397937697
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
Review comment:
Please document:
- input, what would vmap eventually store
- The general algorithm we use(e.g. we propagate the bits backwards into vmap)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on issue #5092: [PASS] dtype
rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on issue #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#issuecomment-603902373
Some more comments, mainly wrt to clarity of the code, test coverage and efficiency concerns. Thanks for bringing in the PR. given that this is critical to a lot of the codebase :) Let us work to polish it to the best state. Thanks @hzfan for good work so far
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen edited a comment on issue #5092: [TIR][PASS]
dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen edited a comment on issue #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#issuecomment-607938301
Thanks @hzfan for keep improving the code and maintaining a high standard. Thanks @yzhliu @Hzfengsy for helpful reviews, this PR is now merged
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r400065401
##########
File path: src/tir/pass/unroll_loop.cc
##########
@@ -160,7 +160,9 @@ class LoopUnroller : public StmtExprMutator {
PrimExpr extent = tir::Simplify(op->extent);
const IntImmNode *v1 = extent.as<IntImmNode>();
int value = -1;
- if (v1 != nullptr) {
+ // integers that do not fit in int32_t are treated as symbolic,
+ // as it's impossible to unroll such large loops
+ if (v1 != nullptr && v1->value <= std::numeric_limits<int>::max()) {
Review comment:
My motivation here is to prevent overflow in the next line (which uses int):
```
value = static_cast<int>(v1->value);
```
IMO it might be fine to use `int` here, as it's consistent with other parts of the pass. What do you think?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395450594
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,326 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
+ bits_ = bits > bits_ ? bits : bits_;
+ StmtExprVisitor::VisitExpr(e);
+ bits_ = tmp;
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vextent_[op->loop_var.as<VarNode>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vextent_[iv->var.as<VarNode>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vextent_.find(op) != vextent_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const IntImmNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ // the narrowed datatype of Var and IntImm
+ std::unordered_map<const PrimExprNode*, DataType> vmap;
+
+ protected:
+ // internal analyzer
+ arith::Analyzer analyzer_;
+
+ private:
+ // the maximum possible bit of the current expression's return dtype
+ int bits_;
+ // the extent of vars to be rewritten
+ std::unordered_map<const VarNode*, DataType> vextent_;
+};
+
+class DataTypeRewriter : public StmtExprMutator {
+ public:
+ Stmt operator()(Stmt s) {
+ visitor_(s);
+ for (auto i = visitor_.vmap.begin(), last = visitor_.vmap.end(); i != last;) {
+ PrimExpr e = GetRef<PrimExpr>(i->first);
+ if (e.dtype() == i->second) {
+ i = visitor_.vmap.erase(i);
+ } else {
+ ++i;
+ }
+ }
+ return VisitStmt(s);
+ }
+
+ Stmt VisitStmt_(const StoreNode* op) final {
+ PrimExpr value = this->VisitExpr(op->value);
+ is_index_ = true;
+ PrimExpr index = this->VisitExpr(op->index);
+ is_index_ = false;
+ Stmt s = StoreNode::make(op->buffer_var,
+ op->value,
+ index,
+ op->predicate);
+ return StmtExprMutator::VisitStmt_(s.as<StoreNode>());
+ }
+
+ Stmt VisitStmt_(const ForNode* op) final {
+ Stmt s = StmtExprMutator::VisitStmt_(op);
+ op = s.as<ForNode>();
+ PrimExpr e = VisitExpr(op->loop_var);
+ Var var = Downcast<Var, PrimExpr>(e);
+ return ForNode::make(var, cast(var.dtype(), op->min), cast(var.dtype(), op->extent),
+ op->for_type, op->device_api, op->body);
+ }
+
+ Stmt VisitStmt_(const AttrStmtNode* op) final {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ Stmt s = StmtExprMutator::VisitStmt_(op);
+ op = s.as<AttrStmtNode>();
+ const IterVarNode* iv = op->node.as<IterVarNode>();
+ PrimExpr e = VisitExpr(iv->var);
+ Var var = Downcast<Var, PrimExpr>(e);
+ if (ivmap_.find(iv) == ivmap_.end()) {
+ ivmap_[iv] = IterVarNode::make(iv->dom, var, iv->iter_type, iv->thread_tag);
+ }
+ return AttrStmtNode::make(
+ ivmap_[iv],
+ op->attr_key,
+ cast(var.dtype(), op->value),
+ op->body);
+ }
+ return StmtExprMutator::VisitStmt_(op);
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) final {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ if (vmap_.find(op) == vmap_.end()) {
+ vmap_[op] = Var(op->name_hint, visitor_.vmap[op]);
+ }
+ return vmap_[op];
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const SizeVarNode* op) final {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ if (vmap_.find(op) == vmap_.end()) {
+ vmap_[op] = SizeVar(op->name_hint, visitor_.vmap[op]);
+ }
+ return vmap_[op];
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) final {
+ is_index_ = true;
+ PrimExpr index = this->VisitExpr(op->index);
+ is_index_ = false;
+ PrimExpr e = LoadNode::make(op->dtype, op->buffer_var, index, op->predicate);
+ return StmtExprMutator::VisitExpr_(e.as<LoadNode>());
+ }
+
+ PrimExpr VisitExpr_(const IntImmNode* op) final {
+ if (is_index_) {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ return IntImm(visitor_.vmap[op], op->value);
+ }
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const CastNode* op) final {
+ if (is_index_ && visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ PrimExpr e = StmtExprMutator::VisitExpr_(op);
+ const CastNode* new_op = e.as<CastNode>();
Review comment:
assert new_op != nullptr
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r396002811
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,363 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
Review comment:
Fixed.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Hzfengsy commented on a change in pull request
#5092: [TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397949397
##########
File path: src/tir/pass/unroll_loop.cc
##########
@@ -160,7 +160,9 @@ class LoopUnroller : public StmtExprMutator {
PrimExpr extent = tir::Simplify(op->extent);
const IntImmNode *v1 = extent.as<IntImmNode>();
int value = -1;
- if (v1 != nullptr) {
+ // integers that do not fit in int32_t are treated as symbolic,
+ // as it's impossible to unroll such large loops
+ if (v1 != nullptr && v1->value <= std::numeric_limits<int>::max()) {
Review comment:
I wonder if we should use `int32_t` here rather than `int`. I'm not sure, but just worry about `int` will represent different types (`int16_t`, `int32_t` or `int64_t`) on different systems and devices.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on issue #5092: [TIR][PASS] dtype
rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on issue #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#issuecomment-606966333
@hzfan Good work. We are in the progress of migrating to the new transform pass manager API. Can you also add a variant of the pass for IRModule and change the testcases to the new style? We can still keep using the old API until we migrated everything itno the new pass style.
reference https://github.com/apache/incubator-tvm/pull/5198
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397936885
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
+ (bound->max_value <= ubound && bound->min_value >= lbound)) {
+ bits = target_bits_;
+ }
+ int tmp = bits > bits_ ? bits : bits_;
+ std::swap(bits_, tmp);
+ StmtExprVisitor::VisitExpr(e);
+ std::swap(bits_, tmp);
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vextent_[op->loop_var.as<VarNode>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vextent_[iv->var.as<VarNode>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vextent_.find(op) != vextent_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
Review comment:
It would be great to add more comment here. e.g. We are taking maximum bits for all the possible Exprs that a Var occurs.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397939075
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
Review comment:
NOTE: the constant int bound here is not necessarily the most efficient for deep nested expressions.
As we are recursively calling const int bound for all sub-expressions of e as well(when we recursively visit). Perhaps we want to add a const int bound with memoization option that allows the analyzer to pass a memo(of each subexpr to the const int bound).
This can be added as a TODO item, but would be great to be resolved in next few weeks cc @yzhliu
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r398275074
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
Review comment:
To clarify, this code seems to indicate that we can rewrite lower bits into higher ones, and I think we do not need this behavior.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397939075
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
Review comment:
NOTE: the constant int bound here is not necessarily the most efficient for deep nested expressions.
As we are recursively calling const int bound for all sub-expressions of e as well(when we recursively visit). Perhaps we want to add a const int bound with memoization option that allows the analyzer to pass a memo(of each subexpr to the const int bound).
We could add it as a TODO item, or directly do it in this PR, but would be great to be resolved in next few weeks cc @yzhliu
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r396002858
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,363 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
Review comment:
Fixed
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r400254439
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
+ (bound->max_value <= ubound && bound->min_value >= lbound)) {
+ bits = target_bits_;
+ }
+ int tmp = bits > bits_ ? bits : bits_;
+ std::swap(bits_, tmp);
+ StmtExprVisitor::VisitExpr(e);
+ std::swap(bits_, tmp);
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vextent_[op->loop_var.as<VarNode>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vextent_[iv->var.as<VarNode>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vextent_.find(op) != vextent_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
Review comment:
Added. Here we ensure that datatype is not promoted.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen merged pull request #5092: [TIR][PASS] dtype
rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen merged pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r400246427
##########
File path: tests/python/unittest/test_tir_pass_narrow_datatype.py
##########
@@ -0,0 +1,128 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+from tvm import te
+
+
+def lower(sch, args):
+ binds = {}
+ arg_list = []
+ for x in args:
+ if isinstance(x, te.tensor.Tensor):
+ buf = tvm.tir.decl_buffer(x.shape, dtype=x.dtype, name=x.name)
+ assert x not in binds
+ binds[x] = buf
+ arg_list.append(buf)
+ else:
+ raise ValueError("args must be Tensor, Buffer or Var")
+ bounds = te.schedule.InferBound(sch)
+ stmt = te.schedule.ScheduleOps(sch, bounds)
+ stmt = tvm.tir.ir_pass.StorageFlatten(stmt, binds, 64, False)
+ stmt = tvm.tir.ir_pass.NarrowDataType(stmt, 32)
Review comment:
Tests added. `test_slice` covers the last case you mentioned.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r398041071
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
Review comment:
Do you have an example? We previously reviewed some of the scenarios, not seeing needs doing so.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395155823
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
+ bits_ = bits > bits_ ? bits : bits_;
+ StmtExprVisitor::VisitExpr(e);
+ bits_ = tmp;
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vset_.insert(op->loop_var.as<Object>());
+ vextent_[op->loop_var.as<Object>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vset_.insert(iv->var.as<Object>());
+ vextent_[iv->var.as<Object>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vset_.find(op) != vset_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const IntImmNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ // the narrowed datatype of Var and IntImm
+ std::unordered_map<const Object*, DataType> vmap;
+
+ protected:
+ // internal analyzer
+ arith::Analyzer analyzer_;
+
+ private:
+ // the maximum bits of all containing expressions
Review comment:
Yes. I changed the comments accordingly.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on issue #5092: [PASS] dtype
rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on issue #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#issuecomment-601284180
@junrushao1994 @ZihengJiang @merrymercy @Hzfengsy would be great if you can also help to take a look
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397935092
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
Review comment:
Shall we rewrite lower bits into higher ones? cc @yzhliu
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397932487
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
Review comment:
You only need to do Downcast<IntImm>(the second argument is deduced automatically)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r396002834
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,363 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
Review comment:
Added a parameter `target_bits_`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on issue #5092: [TIR][PASS] dtype
rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on issue #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#issuecomment-603902755
@spectrometerHBH @FrozenGene @Hzfengsy please also help to take a look
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395833034
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,363 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
Review comment:
Use ConstIntBound, and check the bound
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395732407
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
Review comment:
Agree. I added a brief description about the motivation and how the pass works.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r397936976
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
+ (bound->max_value <= ubound && bound->min_value >= lbound)) {
+ bits = target_bits_;
+ }
+ int tmp = bits > bits_ ? bits : bits_;
+ std::swap(bits_, tmp);
+ StmtExprVisitor::VisitExpr(e);
+ std::swap(bits_, tmp);
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vextent_[op->loop_var.as<VarNode>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vextent_[iv->var.as<VarNode>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vextent_.find(op) != vextent_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
Review comment:
Add more comments about the algorithm
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r400255099
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
Review comment:
Documented.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395701694
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,326 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeRewriter;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = 64;
+ if (e.dtype().bits() <= 32 ||
+ analyzer_.CanProve(e <= max_value(DataType::Int(32)) &&
+ e >= min_value(DataType::Int(32)))) {
+ bits = 32;
+ }
+ int tmp = bits_;
+ bits_ = bits > bits_ ? bits : bits_;
+ StmtExprVisitor::VisitExpr(e);
+ bits_ = tmp;
+ } else {
+ StmtExprVisitor::VisitExpr(e);
+ }
+ }
+
+ void VisitStmt_(const ForNode* op) {
+ analyzer_.Bind(op->loop_var,
+ Range::make_by_min_extent(op->min, op->extent));
+ vextent_[op->loop_var.as<VarNode>()] = op->extent.dtype();
+ return StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ IterVar iv = Downcast<IterVar>(op->node);
+ CHECK_NE(iv->thread_tag.length(), 0U);
+ analyzer_.Bind(iv->var,
+ Range::make_by_min_extent(0, op->value));
+ vextent_[iv->var.as<VarNode>()] = op->value.dtype();
+ StmtExprVisitor::VisitStmt_(op);
+ } else {
+ StmtExprVisitor::VisitStmt_(op);
+ }
+ }
+
+ void VisitExpr_(const ReduceNode* op) {
+ // Setup the domain information before simplification.
+ for (const IterVar& iv : op->axis) {
+ analyzer_.Bind(iv->var, iv->dom);
+ vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
+ }
+ // Recursively call simplification when necessary.
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const VarNode* op) {
+ if (vextent_.find(op) != vextent_.end()) {
+ int bits = std::min(vextent_[op].bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const IntImmNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_int()) {
+ int bits = std::min(op->dtype.bits(), bits_);
+ if (vmap.find(op) == vmap.end()) {
+ vmap[op] = op->dtype.with_bits(bits);
+ } else {
+ vmap[op] = op->dtype.with_bits(std::max(vmap[op].bits(), bits));
+ }
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ // the narrowed datatype of Var and IntImm
+ std::unordered_map<const PrimExprNode*, DataType> vmap;
+
+ protected:
+ // internal analyzer
+ arith::Analyzer analyzer_;
+
+ private:
+ // the maximum possible bit of the current expression's return dtype
+ int bits_;
+ // the extent of vars to be rewritten
+ std::unordered_map<const VarNode*, DataType> vextent_;
+};
+
+class DataTypeRewriter : public StmtExprMutator {
+ public:
+ Stmt operator()(Stmt s) {
+ visitor_(s);
+ for (auto i = visitor_.vmap.begin(), last = visitor_.vmap.end(); i != last;) {
+ PrimExpr e = GetRef<PrimExpr>(i->first);
+ if (e.dtype() == i->second) {
+ i = visitor_.vmap.erase(i);
+ } else {
+ ++i;
+ }
+ }
+ return VisitStmt(s);
+ }
+
+ Stmt VisitStmt_(const StoreNode* op) final {
+ PrimExpr value = this->VisitExpr(op->value);
+ is_index_ = true;
+ PrimExpr index = this->VisitExpr(op->index);
+ is_index_ = false;
+ Stmt s = StoreNode::make(op->buffer_var,
+ op->value,
+ index,
+ op->predicate);
+ return StmtExprMutator::VisitStmt_(s.as<StoreNode>());
+ }
+
+ Stmt VisitStmt_(const ForNode* op) final {
+ Stmt s = StmtExprMutator::VisitStmt_(op);
+ op = s.as<ForNode>();
+ PrimExpr e = VisitExpr(op->loop_var);
+ Var var = Downcast<Var, PrimExpr>(e);
+ return ForNode::make(var, cast(var.dtype(), op->min), cast(var.dtype(), op->extent),
+ op->for_type, op->device_api, op->body);
+ }
+
+ Stmt VisitStmt_(const AttrStmtNode* op) final {
+ if (op->attr_key == attr::thread_extent ||
+ op->attr_key == attr::virtual_thread) {
+ Stmt s = StmtExprMutator::VisitStmt_(op);
+ op = s.as<AttrStmtNode>();
+ const IterVarNode* iv = op->node.as<IterVarNode>();
+ PrimExpr e = VisitExpr(iv->var);
+ Var var = Downcast<Var, PrimExpr>(e);
+ if (ivmap_.find(iv) == ivmap_.end()) {
+ ivmap_[iv] = IterVarNode::make(iv->dom, var, iv->iter_type, iv->thread_tag);
+ }
+ return AttrStmtNode::make(
+ ivmap_[iv],
+ op->attr_key,
+ cast(var.dtype(), op->value),
+ op->body);
+ }
+ return StmtExprMutator::VisitStmt_(op);
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) final {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ if (vmap_.find(op) == vmap_.end()) {
+ vmap_[op] = Var(op->name_hint, visitor_.vmap[op]);
+ }
+ return vmap_[op];
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const SizeVarNode* op) final {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ if (vmap_.find(op) == vmap_.end()) {
+ vmap_[op] = SizeVar(op->name_hint, visitor_.vmap[op]);
+ }
+ return vmap_[op];
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) final {
+ is_index_ = true;
+ PrimExpr index = this->VisitExpr(op->index);
+ is_index_ = false;
+ PrimExpr e = LoadNode::make(op->dtype, op->buffer_var, index, op->predicate);
+ return StmtExprMutator::VisitExpr_(e.as<LoadNode>());
+ }
+
+ PrimExpr VisitExpr_(const IntImmNode* op) final {
+ if (is_index_) {
+ if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ return IntImm(visitor_.vmap[op], op->value);
+ }
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const CastNode* op) final {
+ if (is_index_ && visitor_.vmap.find(op) != visitor_.vmap.end()) {
+ PrimExpr e = StmtExprMutator::VisitExpr_(op);
+ const CastNode* new_op = e.as<CastNode>();
+ return CastNode::make(visitor_.vmap[op], new_op->value);
+ }
+ return StmtExprMutator::VisitExpr_(op);
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) final;
+ PrimExpr VisitExpr_(const SubNode* op) final;
+ PrimExpr VisitExpr_(const MulNode* op) final;
+ PrimExpr VisitExpr_(const DivNode* op) final;
+ PrimExpr VisitExpr_(const ModNode* op) final;
+ PrimExpr VisitExpr_(const FloorDivNode* op) final;
+ PrimExpr VisitExpr_(const FloorModNode* op) final;
+ PrimExpr VisitExpr_(const MinNode* op) final;
+ PrimExpr VisitExpr_(const MaxNode* op) final;
+ PrimExpr VisitExpr_(const EQNode* op) final;
+ PrimExpr VisitExpr_(const NENode* op) final;
+ PrimExpr VisitExpr_(const LTNode* op) final;
+ PrimExpr VisitExpr_(const LENode* op) final;
+ PrimExpr VisitExpr_(const GTNode* op) final;
+ PrimExpr VisitExpr_(const GENode* op) final;
+ PrimExpr VisitExpr_(const CallNode* op) final;
+
+ private:
+ // the internal visitor to deduce the narrowed dtype
+ DataTypeVisitor visitor_;
+ // a map from Var before rewrite to that after rewrite,
+ // ensures one old Var maps to exactly one new Var
+ std::unordered_map<const VarNode*, Var> vmap_;
+ // a map from IterVar before rewrite to that after rewrite,
+ // ensures one old IterVar maps to exactly one new IterVar
+ std::unordered_map<const IterVarNode*, IterVar> ivmap_;
+ bool is_index_{false};
+};
+
+#define DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC) \
+ PrimExpr DataTypeRewriter::VisitExpr_(const OP* op) { \
+ PrimExpr a = this->VisitExpr(op->a); \
+ PrimExpr b = this->VisitExpr(op->b); \
+ if (a.same_as(op->a) && \
+ b.same_as(op->b)) { \
+ return GetRef<PrimExpr>(op); \
+ } else { \
+ return FUNC(a, b); \
+ } \
+ }
+
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(AddNode, operator+)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(SubNode, operator-)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MulNode, operator*)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(DivNode, div)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(ModNode, truncmod)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(FloorDivNode, floordiv)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(FloorModNode, floormod)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MinNode, min)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MaxNode, max)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(EQNode, operator==)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(NENode, operator!=)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LTNode, operator <)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LENode, operator<=)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GTNode, operator >)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GENode, operator>=)
+
+PrimExpr DataTypeRewriter::VisitExpr_(const CallNode* op) {
+ PrimExpr e = StmtExprMutator::VisitExpr_(op);
+ op = e.as<CallNode>();
+ if (op->call_type == CallNode::PureIntrinsic) {
+ if (op->name == intrinsic::tvm_if_then_else) {
+ return if_then_else(op->args[0], op->args[1], op->args[2]);
Review comment:
It invokes `BinaryOpMatchTypes` to match the dtype of its two operands (For example, i32 + i64 gets converted to i64 + i64). Since I rewrite the dtype of some Var, I may break the rule that different operands of one operation must share one dtype.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r400252957
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
+ int64_t ubound = Downcast<IntImm, PrimExpr>(max_value(DataType::Int(target_bits_)))->value;
+ int64_t lbound = Downcast<IntImm, PrimExpr>(min_value(DataType::Int(target_bits_)))->value;
+ if (e.dtype().bits() <= target_bits_ ||
Review comment:
Yes, I think this pass aligns with what you said. It only narrows and never promotes. This code means lower bits fits into higher bits. For example, Consider `e` is `i.i64 <= j.i64`, a bool expression with only 1 bit. So `e` fits into i32 and does not hinder narrowing i to i32.
https://github.com/apache/incubator-tvm/pull/5092/files#diff-98ae729cf00e30cff311ed80b4a25df9R129 ensures dtype promotion does not occcur.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5092:
[TIR][PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5092: [TIR][PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r400266099
##########
File path: src/tir/pass/narrow_datatype.cc
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file narrow_datatype.cc
+ * \brief narrow the datatype of indexing vars
+ */
+
+#include <tvm/tir/ir_pass.h>
+#include <tvm/tir/op.h>
+#include "../../arith/ir_mutator_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace tir {
+
+// This pass narrows indexing expressions (like StoreNode::Index)
+// that trivially fit into i32 to i32. Considering that i32 indices
+// may be more efficient on some backends (while i64 may be more
+// efficient on others, like llvm), we may want this pass when i32
+// indices are more efficient.
+//
+// For Var v, we determine its dtype by examining all the PrimExpr
+// that contains v, denoted by E = {e_0 = v, e_1, e_2, ..., e_k}.
+// If all expressions in E fit into i32, then we think v can be narrowed
+// to i32.
+//
+// To make an indexing expression i32, we must make sure that every
+// component of that expression is of dtype i32. So besides Var, we
+// rewrite the following inside an indexing expression
+// - Var
+// - IntImm
+// - Cast
+//
+// Algorithm:
+// - Use DataTypeVisitor to determine whether a Var can be narrowed or not.
+// - Use DataTypeRewritter to rewrite the components of an indexing expression.
+
+using arith::Analyzer;
+using arith::IRMutatorWithAnalyzer;
+using arith::ConstIntBound;
+
+class DataTypeVisitor final : public StmtExprVisitor {
+ public:
+ explicit DataTypeVisitor(int target_bits)
+ : bits_(target_bits), target_bits_(target_bits) {}
+
+ void VisitExpr(const PrimExpr& e) {
+ if (e.dtype().is_int()) {
+ int bits = max_bits_;
+ ConstIntBound bound = analyzer_.const_int_bound(e);
Review comment:
What about a state flag like `with_memo` in `class ConstIntBoundAnalyzer`? Once it is set, we first look up in the table when a new Expr comes. We can have an unordered_map like `unordered_map<PrimExpr*, Entry>` in `ConstIntBoundAnalyzer` to achieve this.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395451093
##########
File path: src/tir/pass/unroll_loop.cc
##########
@@ -160,7 +160,9 @@ class LoopUnroller : public StmtExprMutator {
PrimExpr extent = tir::Simplify(op->extent);
const IntImmNode *v1 = extent.as<IntImmNode>();
int value = -1;
- if (v1 != nullptr) {
+ // integers that do not fit in int32_t are treated as symbolic,
+ // as it's impossible to unroll such large loops
Review comment:
should we also print a warning here?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5092:
[PASS] dtype rewrite for indexing variables
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5092: [PASS] dtype rewrite for indexing variables
URL: https://github.com/apache/incubator-tvm/pull/5092#discussion_r395442691
##########
File path: src/tir/pass/rewrite_datatype.cc
##########
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file rewrite_datatype.cc
Review comment:
suggest to describe a bit more about the motivation behind, ppl may not know why narrowing if they don't have the context we discussed.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services