[CUDA] Initial support for dynamic shared memory #8466

masahi · 2021-07-14T02:28:31Z

Only one allocation is allowed for now. See the test case for usage.

@tqchen @junrushao1994 @vinx13

src/runtime/rocm/rocm_module.cc

src/runtime/thread_storage_scope.h

src/tir/transforms/split_host_device.cc

tqchen · 2021-07-16T12:06:34Z

src/runtime/thread_storage_scope.h

@@ -96,6 +98,8 @@ struct StorageScope {
        return "global" + tag;
      case StorageRank::kShared:
        return "shared" + tag;
+      case StorageRank::kDynShared:


Let us use the same storage scope as shared, but introduce a different tag. e.g. shared.dyn. The scope rank was used to rank the memory hierachy and needs to also made consistent with the thread rank. In this case dynamic shared memory is a special case of shared memory

I changed dynamic shared memory to be a tagged memory for kShared, but I had to add a check e->scope.tag != ".dyn" in a few places in lower_device_storage_access_info.cc and storage_rewrite.cc to workaround error/segfault.

tqchen · 2021-07-16T12:10:06Z

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])


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";

tqchen · 2021-07-16T12:16:10Z

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);


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.

src/tir/transforms/split_host_device.cc

tqchen · 2021-07-16T12:19:30Z

Thanks @masahi I left some comments, also cc @Hzfengsy @vinx13 so they are also aware

masahi · 2021-07-17T12:15:45Z

src/runtime/file_utils.cc

 }

 bool FunctionInfo::Load(dmlc::Stream* reader) {
  if (!reader->Read(&name)) return false;
  if (!reader->Read(&arg_types)) return false;
  if (!reader->Read(&thread_axis_tags)) return false;
+  if (!reader->Read(&use_dyn_shared_memory)) return false;


@tqchen I didn't find a way to optionally read this attribute for the binary reader.

I think. let us think about another way to do backward compact then.

Since binary backward compatibility is more important to us here.

Here is one solution that also helps future backward compatibility.

Let us rename thread_axis_tags => launch_param_tags, and use dyn_shared_memory as a tag for the launching parameter of dynamic shared memory. We can then change the ThreadAxisConfig => LaunchParamConfig.

Let us try the ^ change and do a test to make sure we can load back old cuda binaries

If we do the rename thread_axis_tags => launch_param_tags, the JSON reader would break backward compat at helper.DeclareField("launch_param_tags", &launch_param_tags); .

Or shall we do the rename but keep the old attribute name for JSON reader/writer? e.g.
helper.DeclareField("thread_axis_tags", &launch_param_tags);.

Or shall we have both new and old attribute names as optional field, e.g.

helper.DeclareOptionalField("launch_param_tags", &launch_param_tags); helper.DeclareOptionalField("thread_axis_tags", &launch_param_tags); // for backward compatibility

Went with the latter solution, and confirmed that an old saved module can be loaded via both json and binary.

We can remove the thread_axis_tags later

tqchen · 2021-07-19T11:53:30Z

src/runtime/file_utils.cc

 }

 bool FunctionInfo::Load(dmlc::Stream* reader) {
  if (!reader->Read(&name)) return false;
  if (!reader->Read(&arg_types)) return false;
  if (!reader->Read(&thread_axis_tags)) return false;
+  if (!reader->Read(&use_dyn_shared_memory)) return false;


Let us try the ^ change and do a test to make sure we can load back old cuda binaries

masahi · 2021-07-20T13:05:23Z

@tqchen Backward compat issue has been addressed.

masahi · 2021-07-22T05:27:20Z

Thanks @tqchen @vinx13 @yzh119 @Hzfengsy

* send dyn shmem size to runtime * add dyn shared storage scope * associate buffer var and its storage scoe in split_host_device * tried NVPTX but failed with INVALID_PTX error * test stub * dynamic shmem reduce working * log2 issue fixed * nvptx working * refactor llvm shmem allocation * make linkage argument * support rocm too * send dyn shmem param to hip runtime * remove alloc map from split_host_device.cc * remove attr::storage_scope from split_host_device * lint fix * formatting * update calling convention doc * minor update to test * remove log * remove kDynShared, dyn.shared -> shared.dyn * support backward compat * update json/binary reader/writer * thread_axis_tags -> launch_param_tags * ThreadAxisConfig -> LaunchParamConfig * remove use_dynamic_shared_memory from FunctionInfo meta data * revert change in test_tir_ir_builder.py * make sure kUseDynamicSharedMemoryTag is the last tag * remove continue * update doc string following name change * more comment update following name change Co-authored-by: masa <masa@pop-os.localdomain> Co-authored-by: Masahiro Masuda <masahi@129@gmail.com>

masahi force-pushed the dyn_shmem branch from 0160c68 to 4530764 Compare July 14, 2021 02:43

masahi marked this pull request as ready for review July 14, 2021 11:23

tqchen assigned vinx13 and junrushao Jul 14, 2021

Hzfengsy reviewed Jul 14, 2021

View reviewed changes

src/runtime/rocm/rocm_module.cc Outdated Show resolved Hide resolved

yzh119 reviewed Jul 14, 2021

View reviewed changes

src/runtime/thread_storage_scope.h Outdated Show resolved Hide resolved

src/tir/transforms/split_host_device.cc Show resolved Hide resolved

vinx13 approved these changes Jul 15, 2021

View reviewed changes

Hzfengsy approved these changes Jul 16, 2021

View reviewed changes

tqchen requested changes Jul 16, 2021

View reviewed changes

tqchen added the status: need update need update based on feedbacks label Jul 16, 2021

masahi force-pushed the dyn_shmem branch from c7b070a to 91033f7 Compare July 17, 2021 12:08

masahi commented Jul 17, 2021

View reviewed changes

tqchen requested changes Jul 19, 2021

View reviewed changes

masahi force-pushed the dyn_shmem branch from 91033f7 to 4450e02 Compare July 20, 2021 12:53

masahi requested review from areusch, comaniac, icemelon, jroesch, junrushao, merrymercy, yzhliu and zhiics as code owners July 20, 2021 12:53

masahi removed request for areusch, yzhliu, jroesch, icemelon and zhiics July 20, 2021 13:16

masa and others added 21 commits July 21, 2021 09:27

refactor llvm shmem allocation

6d9c7c4

make linkage argument

e0fbac2

support rocm too

d781526

send dyn shmem param to hip runtime

933f9c5

remove alloc map from split_host_device.cc

509a8c1

remove attr::storage_scope from split_host_device

0ea0962

lint fix

84666a4

formatting

5bcfacd

update calling convention doc

3ac8401

minor update to test

283e04c

remove log

4389ccb

remove kDynShared, dyn.shared -> shared.dyn

2085bfe

support backward compat

94b1a78

update json/binary reader/writer

64a9d5e

thread_axis_tags -> launch_param_tags

c333cdc

ThreadAxisConfig -> LaunchParamConfig

19ec309

remove use_dynamic_shared_memory from FunctionInfo meta data

8deab0b

revert change in test_tir_ir_builder.py

7cbc700

make sure kUseDynamicSharedMemoryTag is the last tag

1151a52

remove continue

587e5b6

update doc string following name change

ffc138a

masahi force-pushed the dyn_shmem branch from 614b95d to ffc138a Compare July 21, 2021 00:28

more comment update following name change

b8c05a5

tqchen approved these changes Jul 21, 2021

View reviewed changes

masahi merged commit e95f10f into apache:main Jul 22, 2021

masahi mentioned this pull request Jul 28, 2021

[CUDA] Support multiple TIR-level dynamic shared memory allocations #8571

Merged

junrushao mentioned this pull request Nov 1, 2021

Apache TVM v0.8 Release Note Candidate #9416

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[CUDA] Initial support for dynamic shared memory #8466

[CUDA] Initial support for dynamic shared memory #8466

masahi commented Jul 14, 2021 •

edited

Loading

tqchen Jul 16, 2021 •

edited

Loading

masahi Jul 17, 2021 •

edited

Loading

tqchen Jul 16, 2021

masahi Jul 17, 2021

tqchen Jul 16, 2021

masahi Jul 17, 2021

tqchen commented Jul 16, 2021 •

edited

Loading

masahi Jul 17, 2021

tqchen Jul 19, 2021

tqchen Jul 19, 2021

masahi Jul 20, 2021 •

edited

Loading

masahi Jul 20, 2021

tqchen Jul 20, 2021

tqchen Jul 19, 2021

masahi commented Jul 20, 2021

masahi commented Jul 22, 2021

[CUDA] Initial support for dynamic shared memory #8466

[CUDA] Initial support for dynamic shared memory #8466

Conversation

masahi commented Jul 14, 2021 • edited Loading

tqchen Jul 16, 2021 • edited Loading

Choose a reason for hiding this comment

masahi Jul 17, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Jul 16, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

masahi Jul 20, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

masahi commented Jul 20, 2021

masahi commented Jul 22, 2021

masahi commented Jul 14, 2021 •

edited

Loading

tqchen Jul 16, 2021 •

edited

Loading

masahi Jul 17, 2021 •

edited

Loading

tqchen commented Jul 16, 2021 •

edited

Loading

masahi Jul 20, 2021 •

edited

Loading