masahi commented on a change in pull request #8466:
URL: https://github.com/apache/tvm/pull/8466#discussion_r669982717
##########
File path: src/tir/transforms/split_host_device.cc
##########
@@ -89,6 +92,17 @@ class VarUseDefAnalysis : public StmtExprMutator {
Stmt VisitStmt_(const AllocateNode* op) final {
this->HandleDef(op->buffer_var.get());
+ auto storage_scope =
runtime::StorageScope::Create(GetPtrStorageScope(op->buffer_var));
+ if (storage_scope.rank == runtime::StorageRank::kDynShared) {
+ ICHECK_EQ(use_dyn_shmem_, false) << "Only one dynamic shared memory
allocation is allowed.";
Review comment:
I'm planning to support allocating multiple dynamic shmem at TIR level,
like:
```
A_shared = ib.allocate(dtype, A.shape, scope="dyn.shared")
B_shared = ib.allocate(dtype, B.shape, scope="dyn.shared")
```
Under the hood, multiple allocations will be merged into a single allocation
via `storage_rewrite.cc` pass. But I haven't thought about merging multiple
allocations with hetero-dtype like your example.
Maybe for dynamic shmem, we just allocate a byte array and use `reinterpret`
cast for splitting, assuming it is possible with TIR? Probably
`storage_rewrite.cc` does not work for such scenario. Thought? @tqchen
@junrushao1994 @vinx13
Also @yzh119 do you have a use case for such hetero-dtype dynamic shared
memory?
--
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]
