HomePhabricator

fix loop unroll and jam - operand mapping - imperfect nest case

Authored by bondhugula on Aug 28 2019, 11:42 PM.

Description

fix loop unroll and jam - operand mapping - imperfect nest case

  • fix operand mapping while cloning sub-blocks to jam - was incorrect for imperfect nests where def/use was across sub-blocks
  • strengthen/generalize the first test case to cover the previously missed scenario
  • clean up the other cases while on this.

Previously, unroll-jamming the following nest

affine.for %arg0 = 0 to 2048 {
  %0 = alloc() : memref<512x10xf32>
  affine.for %arg1 = 0 to 10 {
    %1 = affine.load %0[%arg0, %arg1] : memref<512x10xf32>
  }
  dealloc %0 : memref<512x10xf32>
}

would yield

%0 = alloc() : memref<512x10xf32>
%1 = affine.apply #map0(%arg0)
%2 = alloc() : memref<512x10xf32>
affine.for %arg1 = 0 to 10 {
  %4 = affine.load %0[%arg0, %arg1] : memref<512x10xf32>
  %5 = affine.apply #map0(%arg0)
  %6 = affine.load %0[%5, %arg1] : memref<512x10xf32>
}
dealloc %0 : memref<512x10xf32>
%3 = affine.apply #map0(%arg0)
dealloc %0 : memref<512x10xf32>

instead of

module {
    affine.for %arg0 = 0 to 2048 step 2 {
      %0 = alloc() : memref<512x10xf32>
      %1 = affine.apply #map0(%arg0)
      %2 = alloc() : memref<512x10xf32>
      affine.for %arg1 = 0 to 10 {
        %4 = affine.load %0[%arg0, %arg1] : memref<512x10xf32>
        %5 = affine.apply #map0(%arg0)
        %6 = affine.load %2[%5, %arg1] : memref<512x10xf32>
      }
      dealloc %0 : memref<512x10xf32>
      %3 = affine.apply #map0(%arg0)
      dealloc %2 : memref<512x10xf32>
    }

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes tensorflow/mlir#98

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/98 from bondhugula:ujam ddbc853f69b5608b3e8ff9b5ac1f6a5a0bb315a4
PiperOrigin-RevId: 266073460

Event Timeline

A. Unique TensorFlower <gardener@tensorflow.org> committed rGbc2a543225ab: fix loop unroll and jam - operand mapping - imperfect nest case (authored by bondhugula).Aug 28 2019, 11:42 PM