PR Category

CINN

PR Types

Improvements

Description

Eliminate serial or parallel (thread/block-bound) loops that have a loop-invariant body.

We can eliminate a serial loop if it satisfies rule 1 and 2:

1. The loop variable is not used in any load/store indices or numeric computation within child schedule blocks. This ensures that the loop writes the same value to the same index in each iteration.
2. The loop doesn't contain any inplace operations, e.g. a[0] = a[0] + b[0]. In the presence of inplace operations, the loop count matters.

We can eliminate a parallel loop if it also satisfies rule 3 and 4:

3. It doesn't write to the local buffer (for thread-bound loop) or shared memory (for block-bound loop), because the index for these storages implicitly includes the thread/block index, thereby violating rule 1.
4. Its child schedule blocks don't have consumers, otherwise consumers in other blocks/threads may read an undefined value.

When a loop can be eliminated, we:

1. For serial loop, we just replace the For node with its body, and replace all use of the loop variable to 0.
2. For parallel loop, we cannot remove the For node because it contains the binding info. Instead, we wrap the loop body in if (loop_var == 0).

Example 1:

serial for (k, 0, 32):
  ScheduleBlock(A):
    A[0] = B[0]
=>
ScheduleBlock(A):
  A[0] = B[0]

Example 2:

binded for (threadIdx.x, 0, 32):
  ScheduleBlock(A):
    A[0] = B[0]  # assume that A is global
=>
binded for (threadIdx.x, 0, 32):
  if threadIdx.x == 0:
    ScheduleBlock(A):
      A[0] = B[0]

Note: We don't remove a binded loop, but wrap its body in threadIdx.x == 0.

Example 3:

serial for (k, 0, 32):
  ScheduleBlock(A):
    A[0] = A[0] + B[0]
=>
DO NOTHING!

Note: Even though k is not used in the body, as ScheduleBlock(A) contains an inplace operation, each iteration updates A, so we cannot eliminate it.

Example 4:

bind for (blockIdx.x, 0, 32):
  ScheduleBlock(A):
    A[0] = 3.14
bind for (blockIdx.x, 0, 32):
  ScheduleBlock(B):
    B[blockIdx.x] = A[0]
=>
DO NOTHING!

Note: If we eliminate the For node of A, only the block at blockIdx.x == 0 can get the correct value of A[0] when computing B.

Example 5:

serial for (k, 0, 32):
  if (k * 256) + threadIdx.x < 1000:
    ScheduleBlock(A):
      A[threadIdx.x] = B[threadIdx.x]
=>
if (0 * 256) + threadIdx.x < 1000:
  ScheduleBlock(A):
    A[threadIdx.x] = B[threadIdx.x]

Note: Although k is used in the if condition, it doesn't affect the index to update, but just limits how many times to update. So rewriting k to 0 doesn't change the semantics.