[TIR] Make lower_warp_memory support extent(threadIdx.x) < warp_size

[roastduck]

Pass lower_warp_memory lowers memory bound to "warp" scope into the warp shuffle intrinsic. Currently, this pass only supports the situation where the extent of threadIdx.x equals to the warp size. However, CUDA's __shfl has a 3rd parameter width to shuffle variables in half (or 1/4, 1/8, 1/16) of a warp. This PR uses this extra parameter to enable Pass lower_warp_memory when the extent of threadIdx.x is less than the warp size.

Changes:

1. Add a 3rd parameter width and a 4th parameter warp_size to TVM intrinsic tvm_warp_shuffle. The 4th parameter warp_size is used to help a Code Generator to decide whether a width is legal. For example, the OpenCL backend dose not support the width parameter, so it has to check whether width == warp_size. Since currently lower_warp_memory is the only pass that utilize tvm_warp_shuffle, this change will not break any dependencies.
2. Code Generators that lowers tvm_warp_shuffle are modified. Currently, the only two affected Code Generators are CUDA and OpenCL.
3. In lower_warp_memory, find the value of width first, and then alter the IR base on width, instead of based on warp_size. Then, it generate the modified tvm_warp_shuffle intrinsic.
4. A test which runs lower_warp_memory with 1/2 warp size is added.

Can @tqchen, @ZihengJiang or @ajtulloch make a review or suggest any other reviewers?

[tqchen]

cc @wpan11nv @ZihengJiang @vinx13 @yongfeng-nv @Hzfengsy

[tqchen]

Thanks @roastduck. I wonder if we can also discuss the alternative abstractions. Right now the abstraction seems to suggest that conceptually the size of the warp is reduced to half(as the shuffle size). However, another way to view it would be to keep the size of the warp to be fixed(32), but support the index access pattern of the subgroups, for example, the canonical form below describes a shuffle in the group of 4

A[wi] = B[(wi/4)*4+ ((wi % 4) +1) %4]

[roastduck]

Thanks @roastduck. I wonder if we can also discuss the alternative abstractions. Right now the abstraction seems to suggest that conceptually the size of the warp is reduced to half(as the shuffle size). However, another way to view it would be to keep the size of the warp to be fixed(32), but support the index access pattern of the subgroups, for example, the canonical form below describes a shuffle in the group of 4

A[wi] = B[(wi/4)*4+ ((wi % 4) +1) %4]

In the alternative approach, __shfl(x, (threadIdx.x + 1) % 4, 4) becomes __shfl(x, threadIdx.y * 4 + (threadIdx.x + 1) % 4) (and threadIdx.z might also be involved). Is my understanding right?

I think the good thing is better compatibility for OpenCL. By this approach, we can support OpenCL using the old 2-parameter intrinsic.

And I think the bad thing comes with CUDA's new shuffle API. __shfl has actually been deprecated by CUDA, and we will have to switch to the new __shfl_sync API. The new API requires an explicit argument mask to specify which threads are active during this shuffle. In the approach of this PR, we will only need to calculate the activeness within a partial (say 1/2, 1/4, etc) warp, which can be calculated from the if nest, given the thread indices of the current partial warp. But in the alternative approach, we will need to calculate the activeness within a whole warp, which means other threads outside the current partial warp will be involved. It may bring a lot of complexity, and even run time overhead when there are dynamic conditions.

To better support both CUDA and OpenCL, maybe we can use both of the approaches.

[tqchen]

Can we still translate to __shfl(x, (threadIdx.x + 1) % 4, 4) in the alternative approach? given that the pattern(wi= warp_index) directly corrresponds to the related pattern

I see, in this case, it would be great to give a bit more discussion and thought about the new programming model.

• Does that mean we are using a smaller "virtual warp"?(which brings the restriction of not being able to use larger amount of the shuffle)

[roastduck]

I see, so you are talking about the programming model instead of the code generation. In the alternative model, a user doesn't cache a buffer to "warp" scope. Instead, a user accesses variables from other threads directly through indices. Have I correctly understood your idea?

What are the advantages of the new model over the current one? You said the current model conceptually reduces a warp to half. Do you mean that it is hard to mix full-warp shuffling and half-warp shuffling for a single buffer? I think this problem can be solved by improving the detection algorithms used in the lower_warp_memory pass, without switching to a new programming model.

If the expressing ability of the two models are equivalent, I prefer the current "bind to warp scope" model, because it is more general and there is no need to introducing a new language feature.

[tqchen]

I actually meant the "bind to warp scope mode". What worries me a bit is that in our original convention, threadIdx.x always means warp index, and the previous comment was about whether or not we should keep threadIdx.x still equal to be warp size, but to detect patttern
A[wi] = B[(wi/4)*4+ ((wi % 4) +1) %4] (where wi = threadIdx.x) for the half shuffle case

[roastduck]

Let me sort out the two approaches.

• In the current approach, the lowering pass directly detects shuffles on threadIdx.x, which has an arbitrary extent. It may be difficult to detect all the complex use cases, for example when there are mixed half-warp and full-warp shuffling.
• In the alternative approach, we assume threadIdx.x == warp size when detecting shuffles. Therefore, either we require users to set the extent of threadIdx.x to the warp size (and then we close this PR), otherwise we may have to add an additional pass to convert threadIdx.x to match the warp size.

[roastduck]

One way to perform a half-warp shuffle while keeping extent(threadIdx.x) == warp size is like this:

Suppose we have extent(threadIdx.x) == 16, we first split threadIdx.y with factor 2, and then fuse that 2 with threadIdx.x, so the extent of threadIdx.x becomes 32.
2. Then we can perform any lowering procedure with the assumption of threadIdx.x == 32.

Either we require users to write such a schedule, otherwise we modify the schedule for them. I think either way is not intuitive enough for users. It may also hinder debugging, and debugging in TVM is already difficult. Note that we have to modify threadIdx.x in more than one scopes, in order to keep the thread index consistent. Here is a more complex example, which is simplified from the algorithm I am currently working on.

// a is shaped (n)
// b is shaped (16)
// c is shaped (n, 16)
// extent of threadIdx.x == 16
for (i.outer = 0; i.outer < n; i.outer += 16) {
    if (i.outer + threadIdx.x < n) {
        a.warp[i.outer + threadIdx.x] = a[i.outer + threadIdx.x]; // (1)
    }
    for (i.inner = 0; i.inner < min(16, n - i.outer); i.inner++) {
        c[i.outer + i.inner, threadIdx.x] += a.warp[i.outer + i.inner] * b[threadIdx.x]; // (2)
    }
}

threadIdx.x in both statement (1) and (2) will be fused to 32, which is a major modification to the schedule. Users may meet difficulties with this schedule.

[tqchen]

Thanks for the great discussions so far, it would be great to have a discussion(perhaps in the forum) about the possible conventions to define th warp, just to clarify further.

• A0: Enforce threadIdx.x == warp_size, and make it the warp index
  • Add a subwarp shuffle detection directly on to the warp shuffle to detect sub-warp shuffle pattern like B[(wi/4)*4+ ((wi % 4) +1) %4] whcih corresponds to a sub-warp shuffle.
• A1: Warp virualization -- Allow threadIdx.x to be at subwarp level. (this PR)

From what I see, A0 might introduce a bit more complexity, but allows a mixture of subwarp and full warp shuffle patterns.

A1 is can be viewed as "virtualizing the warps" by having more virtual warps that acted at the sub-warp level, and use a single warp to simulate them, we cannot mix that with full warp shuffle.

It would be great if we can think a bit about how to describe these different clearly and document them in the comment of the code, and possibly in the future developer docs.

If my understanding is correct, and we can document A1's concept clearly, I can go ahead and merge this PR first, then we do followup discussons.

[roastduck]

I agree with the description of A0 and A1, and now I think there may be an A2 in the design space.

• A2: Improve the shuffle detection. We can analyze all the three thread axes, instead of only threadIdx.x. If the detection is strong enough, we can do whatever shuffle we want without limiting threadIdx.x to be the warp axis.

I think A2 may require less complexity than A0, where we should alter the thread indices.

And for documentation, am I suppose to refine the comments in the code, or start a new documentation page for the topic of warp memory?

[tqchen]

We can start with refining the comments in the code. More docs about warp semantics is always more than welcomed as followup PRs

[roastduck]

Added a bit more descriptions for width. I think it's enough to explain the sub-warp concept.

[wpan11nv]

Are there any real-world usages of this sub-warp scope? Should we upgrade TVM to emit new warp level APIs and then think about extensions?

[tqchen]

@wpan11nv can you try to comment about the new warp API and potential backward compact problem of this extension?

[roastduck]

Can we merge this PR first before discussing about a new API?

[tqchen]

OK, this is merged for now, let us open new RFC threads for the new warp API in CUDA. @roastduck I think @wpan11nv meant the API in CUDA instead of TVM's convention.

Thanks @roastduck @wpan11nv