[Tracking Issue] BufferAccess Migration Followup Items

[tqchen]

This is a tracking issue created for followup tracking items in #9727.

The particular PR contains an important new feature(transform_layout). The main goal of this issue is to do followup shepherding around the IR change part(buffer access and related allocations), which have a broader impact. The particular PR takes many rounds of reviews and but still have quite a few points that would benefit from systematic clarification, cleanup and further polish.

Given that there is an important followup(of layout transform) that depend on the change, and the PR is already at a state where incremental changes can further block the progress as other concurrent changes are made. We took the deliberation to merge PR, then continue to do post-hoc shepherding. This decision is supposed to happen in rare circumstance with clear followup actions to backup the decisions.

This is the tracking issue for followup shepherd actions

High-level gist of change

Load/Store are deprecated and changes to BufferLoad/Store of 1D array. BufferLoad/Store of 1D elements are being used in the place of Load/Store. Overall this is a good directly as we unify more analysis onto the (possibly multi-dimensional) buffer access, opens doors for special memory unit that are in nature multi-dimensional.

However, we do need more clarifications and systematic review/cleanup (see check list).

Check List

• Clarify the buffer access rule, aliasing and allocation rules
  • WIP (RFC: clarifying buffer declaration and access tvm-rfcs#63)
• Buffer.data pointer type and buffer element type relation
• Create a migration guide for people to upgrade from the Load/Store
• Use WithIndices for NormalizeDType
• Find alternative to preflatten_buffer_map (possible via buffer redeclaration)
  • [TIR] Remove PrimFuncNode::preflattened_buffer_map #10940
• Discuss and narrow down the scope of address_of, ideally to 1D acess
• Discuss and clarify the scope of vector buffer indexing rule (RFC: clarifying buffer declaration and access tvm-rfcs#63)
• Post-hoc Revisit key passes involved
  • CodegenLLVM
  • StorageRewrite
  • Vectorize
  • VectorPointerTypeRewrite
  • ThreadStorageSync
  • StorageFlatten
  • MergeDynamicSharedMemoryLocations
  • LowerWarpMemory
  • LowerThreadAllreduce
  • NarrowDataType
  • Other passes that are involved in the refactoring

cc @Hzfengsy @junrushao

[tqchen]

cc @vinx13 @junrushao1994 @masahi @Lunderberg

[masahi]

Can we remove Load / Store node altogether now?

[tqchen]

Let us first finish the items, streamline changes and likely remove load/store in next release. So we have sufficent timewindow and error message for people to react

[Lunderberg]

Discuss and clarify the scope of vector buffer indexing rule

I agree with adding additional restrictions. I've created draft PRs #10512 and #10513, to test the impact of different rules on existing test cases. Each requires that only a single index can have multiple lanes, while #10513 additional requires that to be the last index.

@tqchen After thinking on your comment here, I agree that if there are no issues introduced by using the stricter restriction, we should use it, as we can also loosen the condition later if needed.

[wrongtest-intellif]

Thanks for the great efforts!

Here is my question: Whether we will have some SSA (or no-alias?) property on Buffer object in low-level tir? Previously we can treat every buffer var uniquely defined by Allocate and used by Load/Store/Opaque accesses. Now there seems to be some unclearness. For example, what if we create multiple buffers on the same buffer var associated by Allocate stmt?

[tqchen]

@wrongtest great point. I agree it is important to summarize the aliasing property. @vinx13 perhaps can add more comments on this

[vinx13]

While tir.noalias property still hold for PrimFunc args, T.decl_buffer creates a buffer alias if the underlying data variable (.data field) overlaps with another buffer. Buffer created via alloc_buffer always do not alias. Buffer aliases do not need Allocate to create the data variable. Each data variable can be allocated via Allocate once. If a transformation would produce multiple allocations of the same buffer var (e.g. unrolling a loop that contains an allocation), the transform should update the allocations to be unique using tvm::tir::ConvertSSA

Buffers should not alias each other unless necessary, because aliased buffers increase complexity for TIR transformations. Passes that rewrite buffers should clearly indicate how aliased buffers are handled. For example, when changing the underlying layout of stored elements in a buffer, all buffer aliases must also be updated. While buffer aliasing is typically free at runtime, this imposes a cost for buffer aliasing both to compile times and development complexity.

We should also extend VeryfySSA pass to also take buffer into consideration.