vinx13 commented on a change in pull request #9727:
URL: https://github.com/apache/tvm/pull/9727#discussion_r808551917
##########
File path: src/target/source/codegen_c.cc
##########
@@ -158,78 +158,52 @@ void CodeGenC::PrintSSAAssign(const std::string& target,
const std::string& src,
}
// Print a reference expression to a buffer.
-std::string CodeGenC::GetBufferRef(DataType t, const VarNode* buffer, PrimExpr
index) {
+std::string CodeGenC::GetBufferRef(DataType t, const BufferNode* buffer,
PrimExpr index) {
+ const VarNode* buffer_var = buffer->data.get();
std::ostringstream os;
- std::string vid = GetVarID(buffer);
+ std::string vid = GetVarID(buffer_var);
std::string scope;
- if (alloc_storage_scope_.count(buffer)) {
- scope = alloc_storage_scope_.at(buffer);
+ if (alloc_storage_scope_.count(buffer_var)) {
+ scope = alloc_storage_scope_.at(buffer_var);
}
- bool is_vol = IsVolatile(buffer);
- if (t.lanes() == 1) {
- if (!HandleTypeMatch(buffer, t) || is_vol) {
- os << "((";
- if (is_vol) {
- os << "volatile ";
- }
- // Scope may not be part of type.
- if (!scope.empty() && IsScopePartOfType()) {
- PrintStorageScope(scope, os);
- }
- PrintType(t, os);
- os << "*)" << vid << ')';
- } else {
- os << vid;
- }
- os << "[(";
- PrintExpr(index, os);
- os << ")";
- if (t.bits() == 4 || (t.bits() == 1 && t.is_int())) {
- os << " / " << (32 / t.bits());
- }
- os << ']';
- } else {
Review comment:
does this case go away with the new vectorization impl?
##########
File path: src/tir/transforms/vectorize_loop.cc
##########
@@ -62,30 +62,89 @@ class VecAllocAccess : public StmtExprMutator {
public:
VecAllocAccess(const VarNode* buf, Var var, int var_lanes)
: buf_(buf), var_(var), var_lanes_(var_lanes) {}
- // Load
+
PrimExpr VisitExpr_(const LoadNode* op) final {
- PrimExpr expr = StmtExprMutator::VisitExpr_(op);
- op = expr.as<LoadNode>();
- if (op->buffer_var.get() == buf_) {
- return Load(op->dtype, op->buffer_var, op->index * var_lanes_ + var_,
op->predicate);
- } else {
- return expr;
- }
+ LOG(FATAL) << "Unexpected use of deprecated LoadNode. Please use
BufferLoadNode instead.";
+ return PrimExpr();
}
- // Store
+
Stmt VisitStmt_(const StoreNode* op) final {
- Stmt stmt = StmtExprMutator::VisitStmt_(op);
- op = stmt.as<StoreNode>();
- if (op->buffer_var.get() == buf_) {
- return Store(op->buffer_var, op->value, op->index * var_lanes_ + var_,
op->predicate);
+ LOG(FATAL) << "Unexpected use of deprecated StoreNode. Please use
BufferStoreNode instead.";
+ return Stmt();
+ }
+
+ PrimExpr VisitExpr_(const BufferLoadNode* op) final {
+ auto load = Downcast<BufferLoad>(StmtExprMutator::VisitExpr_(op));
+ return UpdateBufferAccess(load);
+ }
+
+ Stmt VisitStmt_(const BufferStoreNode* op) final {
+ auto store = Downcast<BufferStore>(StmtExprMutator::VisitStmt_(op));
+ return UpdateBufferAccess(store);
+ }
+
+ private:
+ template <typename Node>
+ Node UpdateBufferAccess(Node node) {
+ // Only update the buffer that's being replaced.
+ if (node->buffer->data.get() != buf_) {
+ return node;
+ }
+
+ arith::Analyzer analyzer;
Review comment:
nit: make it a class member
##########
File path:
tests/python/unittest/test_tir_transform_lower_cross_thread_reduction.py
##########
@@ -25,7 +25,14 @@
def _check(original, transformed):
mod = tvm.IRModule.from_expr(original)
mod = tvm.tir.transform.LowerCrossThreadReduction()(mod)
- tvm.ir.assert_structural_equal(mod["main"], transformed, True)
+ try:
+ tvm.ir.assert_structural_equal(mod["main"], transformed, True)
+ except ValueError:
+ with open("temp_expected.txt", "w") as f:
+ f.write(transformed.script())
+ with open("temp_observed.txt", "w") as f:
+ f.write(mod["main"].script())
Review comment:
remove 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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
