tqchen commented on a change in pull request #8466:
URL: https://github.com/apache/tvm/pull/8466#discussion_r671194376
##########
File path: include/tvm/tir/function.h
##########
@@ -240,7 +240,7 @@ namespace attr {
*
* Call(f,
* [arg1, arg2, ..., arg_n,
- * work_size_1, work_size_2, ... work_size_m])
+ * work_size_1, work_size_2, ... work_size_m, dyn_shmem_size])
Review comment:
Let us introduce a special meta-data to indicate that dynamic shared
memory is used. This is to make sure the calling convention is backward
compatible when dyn shared memory is not provided.
```
constexpr const char* kDeviceUseDynSharedMemory =
"tir.device_use_dyn_shared_memory";
```
##########
File path: src/runtime/thread_storage_scope.h
##########
@@ -223,6 +232,7 @@ class ThreadAxisConfig {
w.work_size[arg_index_map_[i]] = size;
}
}
+ w.dyn_shmem_size = static_cast<size_t>(x.values[base_ +
arg_index_map_.size()].v_int64);
Review comment:
We need to consider backward compatibility for kernels where dynamic
shared memory is not used.
Update https://github.com/apache/tvm/blob/main/src/runtime/meta_data.h#L103
to add a `use_dyn_shared_memory`(default to false) field, update json reader
to use DeclareOptionalField, to optionally read this field in. Only decode the
value when the flag is true.
##########
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:
We can do a lowering stage that translate allocate dynamic memory to an
intrinsic that directly get the dynamic memory with an offset(where multiple
allocation can be merged with proper alignment). This will help us get around
some limitation imposed by storage rewrite. Split host device might be the
right place for 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]
