You are viewing a plain text version of this content. The canonical link for it is here.

Posted to dev@tvm.apache.org by Lunderberg <no...@github.com.INVALID> on 2021/10/05 21:15:01 UTC

[apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

This RFC introduces a hard boundary between the “logical layout” of a mathematical tensor and the “physical layout” of a buffer in memory, along with a specification for defining the conversion between the two.
You can view, comment on, or merge this pull request online at:

  https://github.com/apache/tvm-rfcs/pull/39

-- Commit Summary --

  * <a href="https://github.com/apache/tvm-rfcs/pull/39/commits/5e3773ec22dbf44f6a3d771afba7e46d360b9c47">[RFC][TIR] Separate physical and logical layout of buffers</a>

-- File Changes --

    A rfcs/XXXX-buffer-physical-layout.md (399)

-- Patch Links --

https://github.com/apache/tvm-rfcs/pull/39.patch
https://github.com/apache/tvm-rfcs/pull/39.diff

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Wuwei Lin <no...@github.com.INVALID>.

One way to represent the layout mapping in TIR is to introduce different storage scopes and have a registry of pre-defined layout mapping (for example, we already did similar thing for [`wmma` fragments](https://github.com/apache/tvm/blob/813136401a11a49d6c15e6013c34dd822a5c4ff6/python/tvm/topi/cuda/tensor_intrin.py#L94-L101), a special data structure for tensor core input). The cons is that the TIR itself doesn't contain the mapping, the meaning of TIR depends on the contents in the registry. If we can have a few pre-defined mapping that is general to be used in different operators, this might be fine. A possible way to embed the mapping inside TIR is to build the `PrimExpr` from the mapping function, which is similar to what we did for [te.compute](https://github.com/apache/tvm/blob/main/python/tvm/te/operation.py#L106)

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-936668384

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by wrongtest <no...@github.com.INVALID>.

@Lunderberg Hi, I am much interested in `transform_layout` but my team depends totally on TensorIR schedule instead of TE. Could you kindly provide more design points on TensorIR side? It would be great if we can enjoy this preview feature in TensorIR. It is really useful for us.

We have implemented some TensorIR primitives to serve similar purposes in form below to mutate the `Buffer` object's layout, strides and dtypes.
```python
s.replace_buffer_obj(block, write_buffer_idx, *a set of rewrite callbacks)
```
Since generally all buffer accesses are multi-dimensional in TensorIR schedule phase, the implementation is a bit easier (just something like a pass to substitute the buffer object) than in TE, if no extra representative variables are introduced. Is the `transform_layout` would also be like above?
```python
s.transform_layout(block, buffer_idx, remap_func)
```

Another form we use is just a duality of loop transformation primitives, where representative variables for buffer or axis of buffer are tracked into TensorIR schedule states.
```python
n, c, h, w = s.get_buffer_axes(block, buffer_idx)
c_outer, c_inner = s.split_for_buffer(c)
s.reorder_for_buffer(n, c_outer, h, w, c_inner)

# for `transform_layout` look like this?
buffer_rv = s.get_buffer(some identifier)
new_buffer_rv = s.transform_layout(buffer_rv, remap_func)
```
Is it possible to provide both integrated `transform_layout` primitive and step by step primitives for user's convenience?

Very glad to know your opinions! :)

--
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-993336900

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

> I'd suggest adding the BufferTransform data structure here which will be very helpful to other audience.

Sounds good, and I've added a description of it, and a possible data structure for it.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-956454255

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

> How will vectorization work? If there is a `vectorize` directive spanning a logical extent, will the vectorization pass create multidimensional `ramp`s? How will vector loads and stores be represented?

While in principle a vectorized load/store could fallback to a non-vectorized load/store, this would result in ignoring the `vectorize` directive entirely, which would be rather unexpected from the user's perspective.  Therefore, my plan for the initial implementation is to require the elements to be contiguous in both the logical and physical layouts if they are vectorized.  In both the physical and logical layouts, the vector load/store would be represented using a `RampNode` with `stride==1`.

This restriction could then be loosened by allowing schedules to be expressed in terms of the physical layout, rather than the logical layout.  I haven't thought through the full implications of this, which would likely involve either introducing an inlined Stage whose axes are the physical axes, or by exposing the physical axes as a separate parameter within the same stage.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-937824488

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Junru Shao <no...@github.com.INVALID>.

Thank you all for the discussion @Lunderberg @vinx13 @areusch!

-- 
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-1043226453
You are receiving this because you are subscribed to this thread.

Message ID: <ap...@github.com>

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Wuwei Lin <no...@github.com.INVALID>.

Thanks for adding the discussion points.

I understand the difficulty implementing it as eager transform in TE, mainly because most other schedule primitives were not done eagerly as in TIR. So adding a rewrite pass for `BufferTransform` makes sense to me. 

> Should BufferTransform apply only to its body, or apply to the entire graph it is contained in?
> Option 2 is preferred.

When adding `BufferTransform` as a statement, we will need to make sure the semantic is clear. In the provided example, adding `BufferTransform` immediately after `Allocate` has clear meaning of Option2 that it should be applies to everywhere the buffer appeared. What if `BufferTransform` appear at the middle, e.g. `Allocate -> some buffer accesses -> BufferTransform -> some another buffer accesses`. Since `Option2` suggests the transform is global, shall we consider `BufferTransform` being part of function attribute?

> When the transformation is performed, would that include fusing axes together?

My thoughts on TIR and `FlattenBuffer` is that `buffer_transform` will return unfused axes, and `FlattenBuffer` should decide how to fuse them according to buffer scope or other auxiliary information. This implies separation between buffer transform and the underlying physical layout requirement (e.g. 1-D by default in most cases). I think the idea here is close to T2. Using `AXIS_SEPARATOR` is also a potential way to specify the behavior of `FlattenBuffer`. On the other hand, we might also need deliberation on whether buffer transform and flattening to physical buffer should be coupled together.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-959796056

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Wuwei Lin <no...@github.com.INVALID>.

I'd suggest adding the `BufferTransform` data structure here which will be very helpful to other audience.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-955098088

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

> Since Option2 suggests the transform is global, shall we consider BufferTransform being part of function attribute?

I had initially placed `BufferTransform` as a statement so that it could be possible to extended it to have a transformation defined by references to variables within the function body.  This would allow other transforms that require rewriting memory layouts to add a `BufferTransform` defining the layout, relying on a later pass to implement the transformation itself.  For example, `tir.transform.InjectDoubleBuffer` could be implemented by adding a buffer transform `lambda *indices: [loop_iter%2, *indices]`, and the load/store rewrites of `tir.transform.VectorizeLoop` could be implemented by adding a buffer transform `lambda *indices: [*indices, loop_iter]`.

On further thought, this potential future extension would still be possible with the definition in the `PrimFunc`, so long as visitors that modify the `loop_iter` also modify the transformation, so I support having the transformations as a global definition, and will update the RFC.

Having a single global definition of the buffer transformations would also make the order clearer in cases of multiple transformations.  A transformation may be more easily expressed multiple sequential transformations, such as an initial transformation of a matrix into tiles for transfer to a GPU, then another transformation for ordering the elements into groups of size `warp_size`.  If these are in different portions of the TIR call graph, the order of these transformations would depend on the traversal order.

> This implies separation between buffer transform and the underlying physical layout requirement (e.g. 1-D by default in most cases).

That would minimize the number of changes to the TIR, though it would increase the number of distinct changes that are made while lowering.  I had been thinking of the flattening to the underlying physical layout requirement as a special case of a transformation, which happens to define a row-major traversal.

Good point, though, on how the current flattening depends on the parameters outside of the buffer itself.  As such, the transformation representing the flattening couldn't be generated initially.  The concepts could be combined at some point in the future, if the flattening is made to depend only on the buffer, but that should be a separate change.

Between those two points, I think that would be that `PrimFuncAttrs::attrs` should include a map from buffers to a list of layout transformations, and a map from buffers to a list of axis separator locations.  The layout transformations would be consumed by the pass that rearranges the layout, and the list of axis separator locations would be consumed by either `StorageFlatten` or `FlattenBuffer`.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-960254384

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

Thank you for the comments, @vinx13.

> For example, another way is to use a mapping function: (n, c, h, w) -> (n, tir.floordiv(c, 4), h, w, tir.floormod(c, 4)). This would allow arbitrary mapping (we can add more restrictions like requiring affine mapping though, to make analysis easier).

~~Having an arbitrary mapping function was something that I considered as an alternative, and would like to use as part of the internal representation.  However, for the user-facing API, I prefer having a representation that cannot express an illegal internal state, rather than relying on validation.  The arbitrary mapping function could have more than one logical index map to the same physical index (e.g. `(i,j) -> (i+j,)`), which would need to be caught after the fact.~~

Never mind, I'm convinced after going through your point below about CUDA's permuted layout.

> A possible use cases of more complex mapping is [permuted layout](https://github.com/NVIDIA/cutlass/blob/master/media/docs/implicit_gemm_convolution.md#shared-memory-layouts) for shared memory on CUDA.

Thank you for the example, and this sort of example was what I was hoping to get in order to determine if the representation of reorder/splits is sufficiently flexible.  If I'm reading the link correctly, I agree that the physical layout in shared memory cannot be expressed in the proposed notation.  If the global memory is 2-d memory `(i,j)`, the closest to representing physical layout in shared memory would be either `((i,2), (j,4), (j,8), SEP, (i,4))`.  This would maintain the 128-bit vectors `(j,8)`, transpose because `(j,8)` appears in the first physical coordinate, but wouldn't handle the XOR condition needed to have the conflict-free access.  It also couldn't specify that the `(i,2)` in the first physical coordinate would need to take precedence over the `(i,4)` in the second physical coordinate.  In the functional notation, these would be `(i,j) -> (floormod(i,2), (floormod(row,8)) ^ (floordiv(j,8)), floormod(j,8), SEP, row)`, where `row = floordiv(i, 2)`.

I am convinced, while the earlier notation may be useful as a more readable input in some cases, it doesn't cover all use cases.  At most, it could be an alternative form of input, but the function definition mapping from logical indices to physical indices should be the primary representation.

In terms of how to represent this in the TIR graph, I'll have to think a bit on it.  I'm picturing either `BufferNode` holding `Array<Var> index_vars` and `Array<IterSumExpr> physical_layout`, or making a TIR variant of `tvm::relay::FunctionNode`.

> Also, there are related affine analysis infrastructure available, it would be great if we can reuse it for loop analysis and rewriting.

Thank you for mentioning these.  I intend to use these utilities for the implementation, and will add it to the detail section.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-936560130

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Krzysztof Parzyszek <no...@github.com.INVALID>.

How will vectorization work?  If there is a `vectorize` directive spanning a logical extent, will the vectorization pass create multidimensional `ramp`s?  How will vector loads and stores be represented?

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-937152464

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

Sounds good!  I've updated with examples of scheduling with the returned new axes, which work in the implementation posted in [PR#9727](https://github.com/apache/tvm/pull/9727).

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-992608499

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Tejash Shah <no...@github.com.INVALID>.

cc @manupa-arm 

-- 
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-1042145943
You are receiving this because you are subscribed to this thread.

Message ID: <ap...@github.com>

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Junru Shao <no...@github.com.INVALID>.

Merged #39 into main.

-- 
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#event-6092157169
You are receiving this because you are subscribed to this thread.

Message ID: <ap...@github.com>

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

True, and that would allow both for user-defined mappings, and for specifying standard layouts.  I have a bit of concern with using the `scope` parameter to also describe layouts, since in my mind "scope" refers to where the entire buffer is located, while "layout" refers to how individual elements are arranged within the buffer.  Two buffers that have the same layout and different scopes could be `memcpy`-ed between the two scopes without additional rearrangement, and that would be obscured if the `scope` parameter is used to define both the scope and the layout.

I really like that idea of building the TIR representation based on the provided function.  Like with `te.compute`, that gives a convenient way to allow the library to handle the var/primexpr building, so long as the user provides the function to compute the physical indices.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-937008706

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

Following a video chat discussion with @vinx13 , we touched on a number of points, summarized below. Also, we are adding @vinx13 as a co-author on this RFC.

- Are there cases where the flattening in `StorageFlatten`/`FlattenBuffer` should be inferred from buffer properties, rather than explicitly specified by the user? For example, if a buffer has `"texture"` scope, then we know it must be flattened to a 2-d buffer. We concluded that this wouldn't be possible, because the number of resulting dimensions isn't sufficient to define the flattening being applied. For example, if a 4-d buffer is being flattened to 2-d for use in texture memory, the four initial axes `[A, B, C, D]` could be flattened to `[A, fuse(B,C,D)]`, `[fuse(A,B), fuse(C,D)]`, or `[fuse(A,B,C), D]`, without any clear method that is better or worse.

- How will buffer layout transformations be represented in TensorIR schedules? `buffer_transform` will be a primitive transformation in TensorIR, which is eagerly applied on the TensorIR computation.
- In all cases, this would rewrite the buffer shape, and would rewrite loads/stores of that buffer.
- If these loads/stores occur within a series of nested loops that cover all values of the buffer, and have no additional computation (e.g. cache read/write) in the body of these loops, then the loops will be rewritten to be along the transformed axes. can write remainder of schedule in terms of the transformed axes. Otherwise, rewriting the loops would not be well-defined, and will not be done.
- The recommendation for use will be to apply the layout transformations prior to any other scheduling passes that could impact the loop structure, so that rewriting of the loops is possible.

- Should buffer flattening be implemented as a special case of layout transformation? Buffer flattening should remain a separate concept from the layout transforms. Where all other layout transformations can be performed eagerly, and should be before other scheduling passes, buffer flattening must be performed after other scheduling passes. If it were done eagerly, other passes wouldn't have sufficient information about the structure of the buffer.

- Is deprecating Store/Load acceptable, instead using BufferStore/BufferLoad throughout all lowering steps? Yes, as this gives a single uniform way to access buffers, regardless of the lowering step. The one concern is that we should port all existing functionality. For example, the vload/vstore methods in Buffer, which currently return Load/Store respectively, should not be removed, and instead should be updated to return flattened BufferLoad/BufferStore.

- RampNode should be treated as a compiler internal, and shouldn't be easily constructible by users as indices into buffers. The preferred method to represent vectorized access is to have a buffer access within a vectorized loop, then allow `tir.transform.VectorizeLoop` to insert the RampNode. This matches previous behavior, where RampNode could occur in flattened Store/Load, while BufferLoad/BufferStore avoided RampNodes to maintain easy analysis of accessed locations.

- Passes that change buffer dimensionality (e.g. InjectDoubleBuffer) should either be moved before the StorageFlatten/FlattenBuffer, or should be rewritten to instead resize the buffer, rather than changing the dimensionaltiy. The former would require the pass to also update the axis separators to be used when flattening.

--
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-961358086

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Wuwei Lin <no...@github.com.INVALID>.

the api side of `transform_layout` looks good, let's add additional examples of scheduling with returned new axes and get this in

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-988207987

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

> Usage of te.AXIS_SEPARATOR: It seems this is only used in the API side but not in BufferTransform, would be good to get some clarification.

That's correct, the `te.AXIS_SEPARATOR` only appears in the API for the TE schedules, and not in the TIR graph generated from the TE schedule. I've updated the RFC with a paragraph on this distinction.

For the axes returned, I'd want to add a third option.

* T2: Using `te.AXIS_SEPARATOR` to separate groups of axes, e.g. `AA = s[A].transform_layout(lambda i, j: [j//4, te.AXIS_SEPARATOR, i, j%4])`, directly implies fusion of the `i` and `j%4` axes, but exposes each transformed axis for use in the TE schedule. That is, `AA.op.axis` has three elements, whose iteration bounds are `[A.shape[1]//4, A.shape[0], 4]`.

Then my thoughts on each option are below, with T2 as my preferred option.

* T0: This is easy to write in te, but any additional manipulation of the axis (e.g. changing the loop ordering independent of the data layout) would require splitting the returned axis to match the layout transformation. It would also mean that `lambda *indices: indices` isn't the identity function, because it would flatten the buffer as part of that transformation.
* T1: As you mentioned, this would require explicitly writing down the transformation to fuse the axes, which would be error-prone. It would also make the default behavior, where all buffer axes are fused into a single index, be more of a special case.
* T2: Doesn't require the user to specify the transformation to fuse the axes, and allows the corresponding loops to be easily reorganized. The fusion in the default case, where no `transform_layout` is applied, then follows the same pattern, where a `BufferTransform` node is inserted to fuse all remaining axes into `num_axis_separators+1 == 1` axes.

> From TIR's perspective, the same primitive can be introduced to TIR as well. It will be an eager transformation (without needing to create BufferTransform node) that rewrite the both producer/consumer immediately upon calling it. This is required if we want to further scheduling the transformed block. I can follow up implementing it in TIR later

When the transformation is performed, would that include fusing axes together? If it does, then the `FlattenBuffer` pass would need distinguish between a 2-d buffer that should be flattened to a flat 1-d buffer, and a 2-d buffer that has already been flattened from some larger dimension. If it doesn't, then the `FlattenBuffer` pass would need to distinguish between a N-d buffer that should be flattened to a 1-d buffer, and an N-d buffer that should be flattened to a 2-d buffer. In either case, `FlattenBuffer` would need some additional metadata to know how the buffer should be handled.

I like the idea of having the transformation be done eagerly, to avoid needing to pass through additional information, but I'm not quite seeing how it could be done. (I had been trying an initial implementation that way on the TE side.)

--
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-959696021

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Thierry Moreau <no...@github.com.INVALID>.

@vinx13 @junrushao1994 

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-935140712

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Thierry Moreau <no...@github.com.INVALID>.

CC @kparzysz-quic @manupa-arm 

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-935138066

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

Following discussion with @tqchen , this RFC has had significant updates made.  The major change is that instead of extending the capabilities of `Store` and `Load` nodes to support N-d indices, they would instead be removed in favor of keeping `BufferStore` and `BufferLoad` nodes throughout the lowering process.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-954215388

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Wuwei Lin <no...@github.com.INVALID>.

Thanks for updating the RFC. Here are some follow-up thoughts:

Usage of `te.AXIS_SEPARATOR`: It seems this is only used in the API side but not in `BufferTransform`, would be good to get some clarification. Also I could see some tradeoff here that worth discussions:
- T0: using `te.AXIS_SEPARATOR` to separate groups of axes, e.g. `lambda m,n,p,q: [m, n, te.AXIS_SEPARATOR, p, q]`. The benefit is writing `m, n` directly implies fusion of `m` and `n`.
- T1: each element is the returned array of the lambda is an axis. The benefit is that this looks simpler, but when fusion is needed (e.g m, n in the above example), we need to manually write down the affine mapping (m * length_of_n + n) and make sure it is correct.

The new schedule primitive `buffer_transform`:
From TIR's perspective, the same primitive can be introduced to TIR as well. It will be an eager transformation (without needing to create `BufferTransform` node)  that rewrite the both producer/consumer immediately upon calling it. This is required if we want to further scheduling the transformed block. 

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-958106273

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Lunderberg <no...@github.com.INVALID>.

Following a video chat with @csullivan, documenting some of the key points of the conversation.

* Setting the physical layout in a TE-based schedule has two roles.  One is the rewrite the buffer itself, and the other is to define the order of iteration when writing to the buffer.  In the latter use case, the schedule should have access to the physical axes for use in the schedule.
* Setting the physical layout is roughly equivalent to using an inlined `te.compute` to define a re-shaped tensor.  However, using `te.compute` requires the compute definition to change, whereas `set_physical_layout` is entirely contained in the schedule.
* For a given tensor, either the physical axes or the logical axes may be used in scheduling, not both.  We were unable to think of use-cases where you would want to change the physical layout, but still maintain the loop-order of the logical layout.  One idea that came of this conversation was to have `set_physical_layout` act similar to `cache_read` and `cache_write`, which introduce a new stage.
* If a physical layout is defined for a cache stage, then the generated code should perform the layout transformation as part of generating that cache.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-938057599

Re: [apache/tvm-rfcs] [RFC][TIR] Separate physical and logical layout of buffers (#39)

Posted by Wuwei Lin <no...@github.com.INVALID>.

Thanks @Lunderberg for the RFC. Logical-physical mapping is definitely an important feature. I also implemented something similar for warp memory to support tensor core instructions on GPU, I'm happy to collaborate more to get an unified design.
Some preliminary comments:
The current representation of logical-physical layout mapping is to use an array of axis/factor to define how the logical axes are split/reordered/fused to form the physical axes. This works for the case of packed layout like `NCHW4c`, but we might need to think whether this is a generic way to represent the mapping. For example, another way is to use a mapping function: `(n, c, h, w) -> (n, tir.floordiv(c, 4), h, w, tir.floormod(c, 4))`. This would allow arbitrary mapping (we can add more restrictions like requiring affine mapping though, to make analysis easier). A possible use cases of more complex mapping is [permuted layout](https://github.com/NVIDIA/cutlass/blob/master/media/docs/implicit_gemm_convolution.md#shared-memory-layouts) for shared memory on CUDA.
Also, there are related [affine analysis infrastructure](https://github.com/apache/tvm/blob/main/include/tvm/arith/iter_affine_map.h) available, it would be great if we can reuse it for loop analysis and rewriting.

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-935305808

Re: [apache/tvm-rfcs] [RFC][TIR] Layout transformations on buffer access (#39)

Posted by Wuwei Lin <no...@github.com.INVALID>.

@wrongtest I'm working on the TensorIR side and have a draft version of `transform_layout`. The current implementation is 
```void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, bool is_write_index, const IndexMap& index_map);```
It applies the mapping function to all access indices to the target buffer, infer the new shape of the buffer and substitute it with the new buffer. I'll send a PR for that soon

-- 
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/39#issuecomment-993828128