Lunderberg opened a new pull request #9727: URL: https://github.com/apache/tvm/pull/9727
This is an implementation of [RFC#0039](https://github.com/apache/tvm-rfcs/pull/39), which adds layout transformations to both flat and non-flat memory buffers. The commits in this PR are split into 4 sub-sections. Subsection 1 is a refactor that should maintain all existing functionality, while enabling the feature additions in later subsections. Subsections 2-4 are feature additions, each of which starts with a commit adding unit tests for that feature. 1. Replacing Load/Store nodes with BufferLoad/BufferStore nodes throughout the entire lowering flow. Access into flat memory regions, previously represented by Load/Store nodes, are now represented by BufferLoad/BufferStore nodes that access a 1-d memory space. 2. Implement `.transform_layout()`, a scheduling step in TE that alters the in-memory layout of a buffer. 3. Add `te.AXIS_SEPARATOR` to `.transform_layout()` to define groups of transformed axes, where each group is flattened to a single axis. These non-flat memory regions are represented by N-d buffer objects, as enabled in step 1. 4. Expose the transformed axes as the return value of `.transform_layout()`, similar to the existing `.fuse()` and `.split()` functions, to allow additional manipulation by the user. The commits are organized internally by `git merge --no-ff` commits for each sub-section. Unfortunately, it looks like those don't show up very cleanly in github's linear list of commits, but can be seen with `git log --graph --format=format:'%h - %s %d'`, the output of which is below. <details> <summary>Output of <code>git log --graph --format=format:'%h - %s %d'</code></summary> <pre><code> * e017a87c5 - Breakpoint, expose the transformed axes for use in TE scheduling. (HEAD -> physical_layout) |\ | * d89f21479 - [TE] Return transformed iteration variables | * 9de453da0 - [TE] Rewrite loop iteration order | * 69f59e6ca - [TE] Implement te::Transform | * 12f6c2dd0 - [UnitTest] Added tests for loop iteration order. |/ * 9868fd5c7 - Breakpoint, axis separators defined. |\ | * 1bdb548e3 - [TE] Fill BufferNode::axis_separators from StageNode | * 749a3c651 - [TE] Added Stage::set_axis_separators. | * 3e3c04ebe - [TIR] Added BufferNode::axis_separators | * 3317e1d4c - [UnitTest] Test N-d indices exposed to low-level codegen |/ * 0e426b45a - Breakpoint, layout_transform implemented. |\ | * 7ce30a2dc - [TIR] Expose tir.transform.ApplyPhysicalLayout for testing | * a78ec6f52 - [TIR] Added ApplyLayoutTransforms as part of StorageFlatten. | * db0c190a9 - [TIR] Added PrimFunc attribute "layout_transform_map", filled from TE. | * 1d0fe1098 - [TE] Added Stage.transform_layout to the Python TE interface. | * 46e173d9f - [TE] Added Stage::transform_layout to the C++ TE implementation. | * 5c278837f - [TIR] Added IndexMap class. | * bc9c5d242 - [UnitTest] Add unit tests to test physical layout remapping. |/ * d2b2a52f5 - Breakpoint, removed Store/Load nodes from use. |\ | * 7453a8b22 - Added pre_flattened_shape/pre_flattened_stride fields to Buffer. | * 5bd3fd619 - Replace Store/Load with BufferStore/BufferLoad in ir_builder | * 31698f0d5 - Updated Buffer::vstore/vload to return BufferLoad/BufferStore objects. | * 79c6f1ae0 - Updated tvm::address_of() to hold BufferLoad instead of Load. | * 363d8c507 - Replacing Load/Store in codegens. | * 4687caa4a - Replacing Store/Load in lowering/legalization passes. | * b7daea9c0 - Replacing Store/Load in analysis functions | * 8f8eeb40d - Replacing Store/Load in utility passes. | * 347865de6 - Replacing Store/Load in StorageFlatten | * ffd0737a0 - Removing Store/Load from examples | * be1555d91 - Removing Store/Load from optimization passes | * 9debd1f55 - Replacing Store/Load in Stmt/Expr Visitor/Mutator | * 4916d2c87 - [TIR] Added BufferLoadNode::LegalizeDtype |/ * 5557b8c4e - Improve tvmc error message from lazy-loading frontend imports (#9074) (upstream/main, main) </code></pre> </details> -- 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]
