From 4450e021c8a3f5dade605d22f710c7ee5b86c0a8 Mon Sep 17 00:00:00 2001 From: masa Date: Tue, 20 Jul 2021 21:11:05 +0900 Subject: [PATCH] remove use_dynamic_shared_memory from FunctionInfo meta data --- src/runtime/cuda/cuda_module.cc | 8 +++---- src/runtime/file_utils.cc | 8 +++---- src/runtime/meta_data.h | 4 +++- src/runtime/rocm/rocm_module.cc | 8 +++---- src/runtime/thread_storage_scope.h | 10 ++++++--- src/target/build_common.h | 4 +++- tests/python/unittest/test_tir_ir_builder.py | 22 ++++++++++---------- 7 files changed, 33 insertions(+), 31 deletions(-) diff --git a/src/runtime/cuda/cuda_module.cc b/src/runtime/cuda/cuda_module.cc index 046ccc3449bd6..e3820d8ca5692 100644 --- a/src/runtime/cuda/cuda_module.cc +++ b/src/runtime/cuda/cuda_module.cc @@ -153,13 +153,12 @@ class CUDAWrappedFunc { public: // initialize the CUDA function. void Init(CUDAModuleNode* m, ObjectPtr sptr, const std::string& func_name, - size_t num_void_args, const std::vector& launch_param_tags, - bool use_dyn_shared_memory) { + size_t num_void_args, const std::vector& launch_param_tags) { m_ = m; sptr_ = sptr; func_name_ = func_name; std::fill(fcache_.begin(), fcache_.end(), nullptr); - launch_param_config_.Init(num_void_args, launch_param_tags, use_dyn_shared_memory); + launch_param_config_.Init(num_void_args, launch_param_tags); } // invoke the function with void arguments void operator()(TVMArgs args, TVMRetValue* rv, void** void_args) const { @@ -242,8 +241,7 @@ PackedFunc CUDAModuleNode::GetFunction(const std::string& name, if (it == fmap_.end()) return PackedFunc(); const FunctionInfo& info = it->second; CUDAWrappedFunc f; - f.Init(this, sptr_to_self, name, info.arg_types.size(), info.launch_param_tags, - info.use_dyn_shared_memory); + f.Init(this, sptr_to_self, name, info.arg_types.size(), info.launch_param_tags); return PackFuncVoidAddr(f, info.arg_types); } diff --git a/src/runtime/file_utils.cc b/src/runtime/file_utils.cc index ea3f67e94fc2f..35832e83f59cf 100644 --- a/src/runtime/file_utils.cc +++ b/src/runtime/file_utils.cc @@ -44,7 +44,6 @@ void FunctionInfo::Save(dmlc::JSONWriter* writer) const { writer->WriteObjectKeyValue("name", name); writer->WriteObjectKeyValue("arg_types", sarg_types); writer->WriteObjectKeyValue("launch_param_tags", launch_param_tags); - writer->WriteObjectKeyValue("use_dyn_shared_memory", use_dyn_shared_memory); writer->EndObject(); } @@ -53,8 +52,9 @@ void FunctionInfo::Load(dmlc::JSONReader* reader) { std::vector sarg_types; helper.DeclareField("name", &name); helper.DeclareField("arg_types", &sarg_types); - helper.DeclareField("launch_param_tags", &launch_param_tags); - helper.DeclareOptionalField("use_dyn_shared_memory", &use_dyn_shared_memory); + helper.DeclareOptionalField("launch_param_tags", &launch_param_tags); + helper.DeclareOptionalField("thread_axis_tags", + &launch_param_tags); // for backward compatibility helper.ReadAllFields(reader); arg_types.resize(sarg_types.size()); for (size_t i = 0; i < arg_types.size(); ++i) { @@ -66,14 +66,12 @@ void FunctionInfo::Save(dmlc::Stream* writer) const { writer->Write(name); writer->Write(arg_types); writer->Write(launch_param_tags); - writer->Write(use_dyn_shared_memory); } bool FunctionInfo::Load(dmlc::Stream* reader) { if (!reader->Read(&name)) return false; if (!reader->Read(&arg_types)) return false; if (!reader->Read(&launch_param_tags)) return false; - if (!reader->Read(&use_dyn_shared_memory)) return false; return true; } diff --git a/src/runtime/meta_data.h b/src/runtime/meta_data.h index 83c042b667d8c..002012a1e1ccd 100644 --- a/src/runtime/meta_data.h +++ b/src/runtime/meta_data.h @@ -99,12 +99,14 @@ Module MetadataModuleCreate( const std::unordered_map& metadata, const std::unordered_map>& sym_vars); +/*! \brief A tag to specify whether or not dynamic shared memory is used */ +constexpr const char* kUseDynamicSharedMemoryTag = "tir.use_dyn_shared_memory"; + /*! \brief function information needed by device */ struct FunctionInfo { std::string name; std::vector arg_types; std::vector launch_param_tags; - bool use_dyn_shared_memory{false}; void Save(dmlc::JSONWriter* writer) const; void Load(dmlc::JSONReader* reader); diff --git a/src/runtime/rocm/rocm_module.cc b/src/runtime/rocm/rocm_module.cc index affae370d9e33..e02a2c683fdb6 100644 --- a/src/runtime/rocm/rocm_module.cc +++ b/src/runtime/rocm/rocm_module.cc @@ -147,13 +147,12 @@ class ROCMWrappedFunc { public: // initialize the ROCM function. void Init(ROCMModuleNode* m, ObjectPtr sptr, const std::string& func_name, - size_t num_void_args, const std::vector& launch_param_tags, - bool use_dyn_shared_memory) { + size_t num_void_args, const std::vector& launch_param_tags) { m_ = m; sptr_ = sptr; func_name_ = func_name; std::fill(fcache_.begin(), fcache_.end(), nullptr); - launch_param_config_.Init(num_void_args, launch_param_tags, use_dyn_shared_memory); + launch_param_config_.Init(num_void_args, launch_param_tags); } // invoke the function with void arguments void operator()(TVMArgs args, TVMRetValue* rv, void* packed_args, size_t packed_nbytes) const { @@ -197,8 +196,7 @@ PackedFunc ROCMModuleNode::GetFunction(const std::string& name, if (it == fmap_.end()) return PackedFunc(); const FunctionInfo& info = it->second; ROCMWrappedFunc f; - f.Init(this, sptr_to_self, name, info.arg_types.size(), info.launch_param_tags, - info.use_dyn_shared_memory); + f.Init(this, sptr_to_self, name, info.arg_types.size(), info.launch_param_tags); return PackFuncPackedArg(f, info.arg_types); } diff --git a/src/runtime/thread_storage_scope.h b/src/runtime/thread_storage_scope.h index aba25a92729e2..e8bbd87e1006f 100644 --- a/src/runtime/thread_storage_scope.h +++ b/src/runtime/thread_storage_scope.h @@ -29,6 +29,8 @@ #include #include +#include "meta_data.h" + namespace tvm { namespace runtime { @@ -198,13 +200,15 @@ struct ThreadWorkLoad { /*! \brief Thread axis configuration */ class LaunchParamConfig { public: - void Init(size_t base, const std::vector& launch_param_tags, - bool use_dyn_shared_memory = false) { + void Init(size_t base, const std::vector& launch_param_tags) { base_ = base; - use_dyn_shared_memory_ = use_dyn_shared_memory; std::vector filled(6, false); for (size_t i = 0; i < launch_param_tags.size(); ++i) { const std::string& tag = launch_param_tags[i]; + if (tag == kUseDynamicSharedMemoryTag) { + use_dyn_shared_memory_ = true; + continue; + } ThreadScope ts = ThreadScope::Create(tag); arg_index_map_.push_back(ts.rank * 3 + ts.dim_index); filled[ts.rank * 3 + ts.dim_index] = true; diff --git a/src/target/build_common.h b/src/target/build_common.h index e86d13c6e7682..c66c2b52822ec 100644 --- a/src/target/build_common.h +++ b/src/target/build_common.h @@ -57,7 +57,9 @@ inline std::unordered_map ExtractFuncInfo(co } } if (auto opt = f->GetAttr(tir::attr::kDeviceUseDynSharedMemory)) { - info.use_dyn_shared_memory = opt.value(); + if (opt.value()) { + info.launch_param_tags.push_back(runtime::kUseDynamicSharedMemoryTag); + } } auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); fmap[static_cast(global_symbol.value())] = info; diff --git a/tests/python/unittest/test_tir_ir_builder.py b/tests/python/unittest/test_tir_ir_builder.py index 0329134bb3fa0..6874678a38924 100644 --- a/tests/python/unittest/test_tir_ir_builder.py +++ b/tests/python/unittest/test_tir_ir_builder.py @@ -538,8 +538,8 @@ def test_device_ir(A, B): s = te.create_schedule(B.op) def check_target(target): - if not tvm.testing.device_enabled(target): - return + # if not tvm.testing.device_enabled(target): + # return freduce = tvm.build(s, [A, B], target) dev = tvm.device(target, 0) @@ -555,13 +555,13 @@ def check_target(target): if __name__ == "__main__": - test_prefetch() - test_if() - test_for() - test_cpu() - test_gpu() - test_while_vectorize() - test_while_collatz() - test_while_mandel() - test_while_binary_search() + # test_prefetch() + # test_if() + # test_for() + # test_cpu() + # test_gpu() + # test_while_vectorize() + # test_while_collatz() + # test_while_mandel() + # test_while_binary_search() test_dyn_shared()