[CINN] TileBroadcastTactic NHWC layout broadcast support

[Enigmatisms]

NOTE: the same PR is posted again in #71464, since the current PR has some problems with git history. Therefore the current PR is discarded.

PR Category

CINN

PR Types

Improvements

Description

Implement the NHWC layout support for TileBroadcastTactic introduced in #70092.

Introduction

The original TileBroadcastTactic only handles the broadcast of the following form:

   [1, C, 1, 1] => [N, C, H, W].

For NHWC layout broadcast, CINN will fall back to TileFirstGeneralTactic since the last axis is not a broadcast axis. It will have some performance issues, such as using a rather static block size for tensor with different channels, which might produces excessive GMEM load.

The extended TileBroadcastTactic also covers the following case:

[1, 1, 1, C] => [N, H, W, C]

as long as the last axis is a preserved axis (see #70092 for the term definition).

Performance Impact

This tactic extension avoids excessive load in the a simple way: find an appropriate block size K, which satifies:

• C % K is 0
• K is a multiple of 32
• K should be as close to C as possible.
• K is within a certain range, for example: [128, 1024].

The first requirement is the most important one: it will eliminate excessive loads by making the load index invariant to the thread coarsening loop index. Here is an example:

Without this extension, for the input tensor with shape (64, 56, 56, 192), the block size is 256, therefore, with thread coarsening, the thread will load one value from GMEM each loop iteration, totaling 4 loads per tensor (that needs to broadcasting).

// example 192-> (64, 56, 56, 192) tensor broadcast
  for (int32_t thread_loop_i = 0; thread_loop_i < 4; thread_loop_i += 1) {
    float var_0_local = var_0[((((thread_loop_i * 256) + (int)threadIdx.x) + ((int)blockIdx.x * 1024)) % 192)];
    /* ... */
    to_broad_cast[/* ... */] = some_func(var_0_local, /* ... */);
}

With this extension, we only need to load once, since the load index is loop-index invariant. Effectively reducing the number of loads required:

// example 192-> (64, 56, 56, 192) tensor broadcast
  for (int32_t thread_loop_i = 0; thread_loop_i < 4; thread_loop_i += 1) {
    float var_0_local = var_0[threadIdx.x];
    /* ... */
    to_broad_cast[/* ... */] = some_func(var_0_local, /* ... */);
}

Limitation

• For tensors with large C, since the block size can not be too high (otherwise the occupancy will be bad), the load index won't be invariant and the loaded data won't be reused. In this tactic extension, we offered a simple (but not perfect) solution: disabling thread coarsening to reduce register requirement, which will in turn increase occupancy. To solve the problem: register requirement must be dealt with.
• This extension supports broadcast with any shape, as long as the broadcast form is [B, P]. For example: (1, H, 1, C) -> (N, H, W, C). But this is not locally tested and can be potentially erroneous.

Experiment Results

Tested on BatchNorm op (data_layout="NHWC"), the forward broadcast kernel (V100):

shape	max bandwidth %	runtime (us)	bandwidth improvement (%)	runtime decrease (%)
256,64,64,192	92.26	1940	35.56	-26.24
128,72,72,224	92.33	1430	64.43	-39.14
256,48,48,224	92.31	1270	67.29	-40.37
400,72,72,384	92.30	7680	62.10	-38.13

Pcard-89620

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[CLAassistant]

All committers have signed the CLA.