[Refactor] Enforce attaching storage scope to PointerType #8366

masahi · 2021-06-29T10:54:27Z

Right now, storage scope information is scattered across different places. This PR proposes to make storage_scope in PointerType, introduced in #8017, be the single source of truth. This enables each Allocate, Store, Load nodes to directly access its storage scope through its associated buffer variable, rather than via attr::stroage_scope + a map between Var and storage scope as currently done today (for example,

tvm/src/tir/transforms/thread_storage_sync.cc

Line 341 in 8131364

std::unordered_map<const VarNode*, StorageScope> storage_scope_;

).

The main changes are:

In schedule_ops.cc, attach storage scope information to ProduceRealizeNode, rather than creating AttrStmt with tir::attr::realize_scope scope.
In schedule_postproc_to_primfunc.cc, always create a new Buffer using storage scope information in ProduceRealizeNode. This ensures that PointerType associated with the new buffer has always storage scope information attached.
In transform passes and codegen classes, try to use storage scope information in AllocateNode directly, and remove attr::storage_scope and attr::realize_scope when they are no longer needed.

Building on this PR, I'll send the following changes later:

Introduce a new storage scope of dynamic shared memory
Remove remaining use of attr::storage_scope
Remove storage scope in Buffer class

tvm/include/tvm/tir/buffer.h

Lines 70 to 71 in 2cca934

/*! \brief storage scope of the buffer, if other than global */

String scope;

tqchen · 2021-06-30T13:22:11Z

Thanks @masahi . seems there are still testcases that we need to fix

masahi · 2021-07-01T05:53:19Z

@tqchen I've hit an issue in lower_thread_allreduce.cc. There are multiple uses of AttrStmt where a storage scope of a buffer variable is assigned after some complicated logic, like this:

tvm/src/tir/transforms/lower_thread_allreduce.cc

Line 94 in ae58f2c

    
           stmt = AttrStmt(repl->buffer_var, attr::storage_scope, StringImm("shared"), stmt);

So I tried to replace this line by something like SetStorageScope(repl->buffer_var, "shared") but this doesn't work because I cannot get a mutable pointer to PointerType via buffer_var->type_annotation.as<PointerTypeNode>().

Creating a new variable also doesn't work because the old buffer_var is shared with other stmt nodes, so updating the buffer_var of AllocateNode results in the use of undefined variable in other nodes.

What should we do about this?

tqchen · 2021-07-01T14:05:04Z

In this particular case, we might need to create a new buffer var and remap the use of the old buffer vars to new one

Jenkinsfile

src/target/llvm/codegen_llvm.cc

src/tir/ir/buffer.cc

src/target/llvm/codegen_llvm.h

This reverts commit b66c3ba.

This reverts commit e650b6c.

This reverts commit 8e20003.

… removal

masahi · 2021-07-07T22:06:09Z

Should we remove scope from Buffer? The scope is now available from the type annotation of the data member in BufferNode.

Yes this is one of my follow-up items. Apparently scope in Buffer is only used by TensorIR related code.

masahi · 2021-07-08T03:44:08Z

src/tir/transforms/update_pointer_storage_scope.h

+namespace tvm {
+namespace tir {
+
+class UpdatePointerStorageScope : public StmtExprMutator {


This is used by passes that require modifying storage scopes after buffer creation. Currently used by lower_thread_allreduce.cc and lower_warp_memory.cc.

masahi · 2021-07-08T03:47:03Z

python/tvm/script/special_stmt.py

@@ -491,6 +491,22 @@ def var(dtype, span):
        super().__init__(var, def_symbol=True)


+@register
+class BufferVarDef(SpecialStmt):


cc @spectrometerHBH for tvmscript changes. This is for distinguishing a normal variable from a buffer variable. The latter one requires storage scope. See the change in test_tvmscript_roundtrip.py. The tvm script printer was also updated.

masahi · 2021-07-08T03:50:06Z

All tests passed, ready for review.

src/tir/transforms/update_pointer_storage_scope.cc

src/runtime/thread_storage_scope.h

src/tir/transforms/update_pointer_storage_scope.cc

src/te/operation/cross_thread_reduction.cc

vinx13 · 2021-07-10T18:20:59Z

src/te/operation/cross_thread_reduction.cc

@@ -177,7 +178,8 @@ Stmt MakeCrossThreadReduction(const ComputeOpNode* self, const Stage& stage,
  std::vector<Var> res_handles(size);
  for (size_t idx = 0; idx < size; ++idx) {
    DataType dtype = reduces[idx]->dtype;
-    res_handles[idx] = Var("reduce_temp" + std::to_string(idx), PointerType(PrimType(dtype)));
+    res_handles[idx] =


shall we remove storage_scope attr here

tvm/src/te/operation/cross_thread_reduction.cc

Line 229 in f459dbf

body = AttrStmt(res_handles[idx - 1], tir::attr::storage_scope, StringImm("local"), body);

tvm/src/te/operation/cross_thread_reduction.cc

Line 234 in f459dbf

AttrStmt(normal_res_handles[idx - 1], tir::attr::storage_scope, StringImm("local"), body);

Removal of AttrStmt creation with attr::storage_scope will come after this PR, where all attr::storage_scope usages will be removed at once. A lot of places in the code base still assume the existence of AttrStmt with attr::storage_scope, so I don't want to remove some AttrStmt while leaving others in this PR.

vinx13 · 2021-07-10T18:22:50Z

src/tir/transforms/lower_thread_allreduce.cc


 namespace tvm {
 namespace tir {

+class UpdatePointerStorageScopeAllReduce final : public UpdatePointerStorageScope {


Can we merge this pass with UpdatePointerStorageScope by adding check for attr::volatile_scope in VisitStmt_(const AttrStmtNode*)

This attr::volatile_scope is created inside this AllocateNode visitor. A visitor for AttrStmtNode don't get to see this.

kparzysz-quic

LGTM

csullivan

@masahi Thanks for the great and healthy refactor!

csullivan · 2021-07-12T18:02:37Z

src/tir/transforms/update_pointer_storage_scope.cc

@@ -0,0 +1,91 @@
+/*


I think we need to also update lower_device_storage_access_info.cc to use the PointerType storage scope in this PR as well.

With a custom storage scope used in a cache_read stage, e.g.

s.cache_read(data, "global.texture", [OL])

I see that lower_device_storage_access_info.cc#L72 is unhappy with the scope tag.

Check failed: (e.info.defined()) is false: Cannot find memory info of global.texture

and this ICHECK failure does not occur on main.

Hmm interesting. attr::storage_scope and the PointerType should have the same storage scope, by the invariant at

tvm/src/tir/ir/stmt.cc

Lines 64 to 73 in 3e5e0ef

if (attr_key == attr::storage_scope) {

const VarNode* buf = node.as<VarNode>();

ICHECK(buf);

const auto* ptr_type = buf->type_annotation.as<PointerTypeNode>();

ICHECK(ptr_type) << "The provided variable is not of pointer type";

auto attr_scope = value.as<StringImmNode>()->value;

ICHECK_EQ(attr_scope, ptr_type->storage_scope)

<< "Storage scopes attached to AttrStmt and buffer var are different. " << attr_scope

<< ", " << ptr_type->storage_scope;

}

.

A possible explanation is that the value of scope.tag.length() at https://github.com/apache/tvm/blob/main/src/tir/transforms/lower_device_storage_access_info.cc#L70 is different between main and this branch. In this case, GetMemoryInfo assumes that a packed function tvm.info.mem.global.texture is defined, but I don't find it anywhere.

Do you have a repro script?

Thanks @masahi. I tracked this issue down to the buffer storage scope being set to global and removing the scope tag. Previously this was fine as the AttrStmtNode realize_scope annotation was used during texture lowering.

This demonstrate the value in removing the buffer storage scope in favor of consolidating around the PointerType's storage_scope. I look forward to that PR! Thanks again for this change.

tqchen · 2021-07-12T18:08:30Z

will let @vinx13 manage this PR

csullivan

LGTM

vinx13 · 2021-07-13T19:33:21Z

Thanks @masahi @kparzysz-quic @csullivan @tqchen this is merged

* Add storage scope to ProducerRealize, always create a buffer with scope * update schedule_ops.cc * update schedule_postproc_to_primfunc.cc * restore more realize_scope This reverts commit b66c3ba. * make the default scope be "" instead of None in ir builder * restore realize_scope visit in storage_flatten.cc * update storage_access.cc * make sure buffer var is of PointerType in ir builder This reverts commit e650b6c. * enforce default storage scope of global * added remap pass but does not work yet * fixed all reduce issue This reverts commit 8e20003. * simplify * trying mitigation for aot test * merge remaining changes from initial branch * remove use of attr::storage_scope from codegen * restore a visit to AttrStmt with attr::storage_scope in storage_rewrite * disable check * lint fix * revert default scope to "" * format * fix volatile access to shared mem in lower all reduce * fixed gpu coorporative load/store test * pass storage scope to PointerType in tvm script parser This reverts commit 99cfb9d18781dcfdea169d920450f9063ab18b6b. * fixed tvmscript roundtrip test * fixed tir flatten buffer test * fixed test_tir_transform_hoist_if.py * use storage scope global by default in aot_executor_codegen.cc * add missing default storage scope in create_primfunc.cc * restore StorageInfo struct in llvm backend * UpdateStorageScope -> WithStorageScope * fixed lower warp memory test * GetStorageScope -> GetPtrStorageScope * Enable storage scope invariant check in AttrStmt constructor * remove GetPtrStorageScope and WithStorageScope from public header * move RemapStorageScope to its own file * add more method to RemapStorageScope * update lower_thread_allreduce to use RemapStorageScope * RemapStorageScope -> UpdatePointerStorageScope * remove realize_scope from hybrid script * removed realize_scope in schedule_ops * remove realize_scope from schedule_postproc_to_primfunc * remove remaining realize_scope usage from schedule_ops.cc * remove realize_scope usage from storage_flatten.cc * fixed test_tir_transform_lower_warp_memory.py following realize_scope removal * Add storage scope to ProducerRealize, always create a buffer with scope * update schedule_ops.cc * update schedule_postproc_to_primfunc.cc * restore more realize_scope This reverts commit b66c3ba. * make the default scope be "" instead of None in ir builder * restore realize_scope visit in storage_flatten.cc * update storage_access.cc * make sure buffer var is of PointerType in ir builder This reverts commit e650b6c. * enforce default storage scope of global * added remap pass but does not work yet * fixed all reduce issue This reverts commit 8e20003. * simplify * trying mitigation for aot test * merge remaining changes from initial branch * remove use of attr::storage_scope from codegen * restore a visit to AttrStmt with attr::storage_scope in storage_rewrite * disable check * lint fix * revert default scope to "" * format * fix volatile access to shared mem in lower all reduce * fixed gpu coorporative load/store test * pass storage scope to PointerType in tvm script parser This reverts commit 99cfb9d18781dcfdea169d920450f9063ab18b6b. * fixed tvmscript roundtrip test * fixed tir flatten buffer test * fixed test_tir_transform_hoist_if.py * use storage scope global by default in aot_executor_codegen.cc * add missing default storage scope in create_primfunc.cc * restore StorageInfo struct in llvm backend * UpdateStorageScope -> WithStorageScope * fixed lower warp memory test * GetStorageScope -> GetPtrStorageScope * Enable storage scope invariant check in AttrStmt constructor * remove GetPtrStorageScope and WithStorageScope from public header * move RemapStorageScope to its own file * add more method to RemapStorageScope * update lower_thread_allreduce to use RemapStorageScope * RemapStorageScope -> UpdatePointerStorageScope * remove realize_scope from hybrid script * removed realize_scope in schedule_ops * remove realize_scope from schedule_postproc_to_primfunc * remove remaining realize_scope usage from schedule_ops.cc * remove realize_scope usage from storage_flatten.cc * fixed test_tir_transform_lower_warp_memory.py following realize_scope removal * Address comments * Remove blank line diff Co-authored-by: Masahiro Masuda <masahi@[email protected]> Co-authored-by: masa <[email protected]>

masahi marked this pull request as draft June 29, 2021 10:54

masahi force-pushed the storage_scope_refactor branch 3 times, most recently from 0c9a912 to 79c6e8e Compare June 29, 2021 20:24

tqchen self-assigned this Jun 30, 2021

tqchen assigned kparzysz-quic Jun 30, 2021

masahi force-pushed the storage_scope_refactor branch from 29ae80c to c443506 Compare June 30, 2021 23:37

masahi force-pushed the storage_scope_refactor branch 4 times, most recently from 553402d to 73c6496 Compare July 5, 2021 10:17

tqchen assigned vinx13 Jul 6, 2021

tqchen reviewed Jul 6, 2021

View reviewed changes

Jenkinsfile Outdated Show resolved Hide resolved

src/target/llvm/codegen_llvm.cc Show resolved Hide resolved

src/tir/ir/buffer.cc Outdated Show resolved Hide resolved

src/tir/ir/buffer.cc Outdated Show resolved Hide resolved

tqchen reviewed Jul 6, 2021

View reviewed changes

src/target/llvm/codegen_llvm.h Show resolved Hide resolved

masahi force-pushed the storage_scope_refactor branch 2 times, most recently from 9f37d55 to ba670d3 Compare July 7, 2021 00:37

Masahiro Masuda and others added 11 commits July 7, 2021 12:17

Add storage scope to ProducerRealize, always create a buffer with scope

a84db69

update schedule_ops.cc

0140631

update schedule_postproc_to_primfunc.cc

b6d8e6c

restore more realize_scope

bebcc50

This reverts commit b66c3ba.

make the default scope be "" instead of None in ir builder

edeaed2

restore realize_scope visit in storage_flatten.cc

cd6167e

update storage_access.cc

a33fb0d

make sure buffer var is of PointerType in ir builder

e878eae

This reverts commit e650b6c.

enforce default storage scope of global

9f98dea

added remap pass but does not work yet

2923f4d

fixed all reduce issue

6a79354

This reverts commit 8e20003.

masahi and others added 3 commits July 8, 2021 07:04

remove remaining realize_scope usage from schedule_ops.cc

086d891

remove realize_scope usage from storage_flatten.cc

e60ad8d

fixed test_tir_transform_lower_warp_memory.py following realize_scope…

cb697d8

… removal

masahi force-pushed the storage_scope_refactor branch from 426b573 to cb697d8 Compare July 7, 2021 22:04

masahi commented Jul 8, 2021

View reviewed changes

minor fix

6a11ab4

kparzysz-quic reviewed Jul 9, 2021

View reviewed changes

src/tir/transforms/update_pointer_storage_scope.cc Outdated Show resolved Hide resolved

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

src/tir/transforms/update_pointer_storage_scope.cc Outdated Show resolved Hide resolved

Address comments

f459dbf

vinx13 reviewed Jul 10, 2021

View reviewed changes

Remove blank line diff

3e5e0ef

kparzysz-quic approved these changes Jul 12, 2021

View reviewed changes

csullivan requested changes Jul 12, 2021

View reviewed changes

tqchen removed their assignment Jul 12, 2021

csullivan approved these changes Jul 12, 2021

View reviewed changes

Merge remote-tracking branch 'upstream/main' into storage_scope_refactor

5710513

vinx13 approved these changes Jul 13, 2021

View reviewed changes

vinx13 merged commit 1a26733 into apache:main Jul 13, 2021

vinx13 added the status: accepted label Jul 13, 2021

This was referenced Jul 14, 2021

[Refactor] Remove scope attribute from Buffer class #8463

Merged

[Refactor] Remove AttrStmt with storage_scope key #8516

Merged

masahi mentioned this pull request Sep 17, 2021

[Bug] BuiltinLower has hard-coded heuristic for alloca which is not appropriate for all kDLCPU target devices #9022

Closed

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

[Refactor] Enforce attaching storage scope to PointerType #8366

[Refactor] Enforce attaching storage scope to PointerType #8366

masahi commented Jun 29, 2021 •

edited

Loading

tqchen commented Jun 30, 2021

masahi commented Jul 1, 2021

tqchen commented Jul 1, 2021

masahi commented Jul 7, 2021

masahi Jul 8, 2021

masahi Jul 8, 2021

masahi commented Jul 8, 2021

vinx13 Jul 10, 2021

masahi Jul 12, 2021

vinx13 Jul 10, 2021

masahi Jul 12, 2021

kparzysz-quic left a comment

csullivan left a comment

csullivan Jul 12, 2021 •

edited

Loading

masahi Jul 12, 2021

csullivan Jul 12, 2021

tqchen commented Jul 12, 2021

csullivan left a comment

vinx13 commented Jul 13, 2021

	/! \brief storage scope of the buffer, if other than global /
	String scope;

	if (attr_key == attr::storage_scope) {
	const VarNode* buf = node.as<VarNode>();
	ICHECK(buf);
	const auto* ptr_type = buf->type_annotation.as<PointerTypeNode>();
	ICHECK(ptr_type) << "The provided variable is not of pointer type";
	auto attr_scope = value.as<StringImmNode>()->value;
	ICHECK_EQ(attr_scope, ptr_type->storage_scope)
	<< "Storage scopes attached to AttrStmt and buffer var are different. " << attr_scope
	<< ", " << ptr_type->storage_scope;
	}

[Refactor] Enforce attaching storage scope to PointerType #8366

[Refactor] Enforce attaching storage scope to PointerType #8366

Conversation

masahi commented Jun 29, 2021 • edited Loading

tqchen commented Jun 30, 2021

masahi commented Jul 1, 2021

tqchen commented Jul 1, 2021

masahi commented Jul 7, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

masahi commented Jul 8, 2021

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

kparzysz-quic left a comment

Choose a reason for hiding this comment

csullivan left a comment

Choose a reason for hiding this comment

csullivan Jul 12, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Jul 12, 2021

csullivan left a comment

Choose a reason for hiding this comment

vinx13 commented Jul 13, 2021

masahi commented Jun 29, 2021 •

edited

Loading

csullivan Jul 12, 2021 •

edited

Loading