Andi Drebes c367a4b6fd feat(compiler): Add batching pass ...

This adds a new pass that is able to hoist operations implementing the
`BatchableOpInterface` out of a loop nest that applies the operation
to the elements of a tensor indexed by the loop induction variables.

Example:

  scf.for %i = c0 to %cN step %c1 {
    scf.for %j = c0 to %cM step %c1 {
      scf.for %k = c0 to %cK step %c1 {
        %s = tensor.extract %T[%i, %j, %k]
        %res = batchable_op %s
        ...
      }
    }
  }

is replaced with:

  %batchedSlice = tensor.extract_slice
       %T[%c0, %c0, %c0] [%cN, %cM, %cK] [%c1, %c1, %c1]
  %flatSlice = tensor.collapse_shape %batchedSlice
  %resTFlat = batchedOp %flatSlice
  %resT = tensor.expand_shape %resTFlat

  scf.for %i = c0 to %cN step %c1 {
    scf.for %j = c0 to %cM step %c1 {
      scf.for %k = c0 to %cK step %c1 {
        %res = tensor.extract %resT[%i, %j, %k]
        ...
      }
    }
  }

Every index of the tensor with the input values may be a quasi-affine
expression on a single loop induction variable, as long as the
difference between the results of the expression for any two
consecutive values of the referenced loop induction variable is
constant.

2022-11-18 12:06:07 +01:00

33 KiB

Raw Blame History

View Raw