[GitHub] [tvm] kparzysz-quic commented on a diff in pull request #10940: [TIR] Remove PrimFuncNode::preflattened_buffer_map

[GitBox <gi...@apache.org>]

kparzysz-quic commented on code in PR #10940:
URL: https://github.com/apache/tvm/pull/10940#discussion_r1038321051


##########
src/tir/transforms/storage_flatten.cc:
##########
@@ -1160,18 +1144,70 @@ class BufferBindUnwrapper : public StmtExprMutator {
   };
 
   const BufferEntry& GetBufferEntry(Buffer buffer) {
-    auto alloc_key = buffer->data.get();
-    if (!buf_map_.count(buffer.get()) && buffer_var_defines_.count(alloc_key)) {
+    if (buf_map_.count(buffer.get())) {
+      const BufferEntry& e = buf_map_[buffer.get()];
+      ICHECK(e.in_scope) << "Cannot access a buffer " << buffer->name << ", out of scope";
+      return e;
+    } else if (buffer_var_defines_.count(buffer->data.get())) {
+      // The buffer var was defined, but the buffer hasn't been seen
+      // before.
       BufferEntry entry;
       entry.buffer = buffer;
+      var_to_buffer_[buffer->data.get()] = buffer;
       buf_map_[buffer.get()] = std::move(entry);
-    }
+      return buf_map_[buffer.get()];
+    } else if (var_remap_.count(buffer->data.get())) {
+      // The buffer var is an alias of a bound buffer.  Only
+      // supported if the bound buffer has no offsets.  In this
+      // case, we just need to make a new aliasing buffer that
+      // shares the remapped data variable.
+      Var old_var = buffer->data;
+      Var new_var = Downcast<Var>(var_remap_[old_var.get()]);
 
-    auto it = buf_map_.find(buffer.get());
-    ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << buffer;
-    const BufferEntry& e = it->second;
-    ICHECK(e.in_scope) << "Cannot access a buffer " << buffer->name << ", out of scope";
-    return it->second;
+      {
+        ICHECK(var_to_buffer_.count(old_var.get()))
+            << "Cannot find remap information for aliased buffer var " << old_var->name_hint
+            << ", required to verify this alias is legal.";
+        const Buffer& aliased_buffer = var_to_buffer_[old_var.get()];
+        const BufferEntry& entry = buf_map_[aliased_buffer.get()];
+        if (entry.remap) {
+          for (const auto& begin : entry.remap->begins) {
+            ICHECK(is_zero(begin)) << "Aliasing of buffer with offset is not supported";
+          }
+        }
+      }
+
+      {
+        Buffer new_buf = buffer;
+        new_buf.CopyOnWrite()->data = new_var;
+
+        RemapInfo remap_info;
+        remap_info.target = new_buf;
+        remap_info.begins = Array<PrimExpr>(buffer->shape.size(), 0);
+        remap_info.extents = buffer->shape;
+
+        BufferEntry entry;
+        entry.buffer = buffer;
+        entry.remap = std::make_unique<RemapInfo>(remap_info);
+        entry.in_scope = true;
+        var_to_buffer_[buffer->data.get()] = buffer;
+        buf_map_[buffer.get()] = std::move(entry);
+      }
+      return buf_map_[buffer.get()];
+    } else if (var_to_buffer_.count(buffer->data.get())) {
+      // This buffer is an alias of a known buffer, with no remaps.  A
+      // buffer entry should be generated and returned.
+      BufferEntry entry;
+      entry.buffer = buffer;
+      entry.in_scope = true;
+      var_to_buffer_[buffer->data.get()] = buffer;
+      buf_map_[buffer.get()] = std::move(entry);
+
+      return buf_map_[buffer.get()];
+    } else {
+      LOG(FATAL) << "Can't work around the undefined buffer";
+      return *static_cast<BufferEntry*>(nullptr);

Review Comment:
   This is not good, please use `__builtin_unreachable()` or something like that...



-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org