• Diego Caballero's avatar
    [MLIR] Enable affine dialect loop fusion (#3290) · aedd8c2e
    Diego Caballero authored
    * [MLIR] Enable affine dialect loop fusion
    
    Enable affine dialect loop fusion in nGraph pipeline. It also adds an
    opt flag to enable/disable it when ngraph-opt is in place. Fusion seems
    to work for simple cases. It wasn't able to fuse dot + add, though, at
    least in my test case. One example that worked:
    
    Input:
      %6 = alloc() : memref<2500x2500xf32>
      affine.for %i3 = 0 to 2500 {
        affine.for %i4 = 0 to 2500 {
          %7 = load %arg0[%i3, %i4] : memref<2500x2500xf32>
          %8 = load %0[%i3, %i4] : memref<2500x2500xf32>
          %9 = addf %8, %7 : f32
          store %9, %6[%i3, %i4] : memref<2500x2500xf32>
        }
      }
      %10 = alloc() : memref<2500x2500xf32>
      affine.for %i5 = 0 to 2500 {
        affine.for %i6 = 0 to 2500 {
          %11 = load %arg2[%i5, %i6] : memref<2500x2500xf32>
          %12 = load %0[%i5, %i6] : memref<2500x2500xf32>
          %13 = addf %12, %11 : f32
          store %13, %10[%i5, %i6] : memref<2500x2500xf32>
        }
      }
      %14 = alloc() : memref<2500x2500xf32>
      affine.for %i7 = 0 to 2500 {
        affine.for %i8 = 0 to 2500 {
          %15 = load %10[%i7, %i8] : memref<2500x2500xf32>
          %16 = load %6[%i7, %i8] : memref<2500x2500xf32>
          %17 = addf %16, %15 : f32
          store %17, %14[%i7, %i8] : memref<2500x2500xf32>
        }
      }
    
    Output:
      %8 = alloc() : memref<2500x2500xf32>
      affine.for %i3 = 0 to 2500 {
        affine.for %i4 = 0 to 2500 {
          %9 = load %arg2[%i3, %i4] : memref<2500x2500xf32>
          %10 = load %2[%i3, %i4] : memref<2500x2500xf32>
          %11 = addf %10, %9 : f32
          %12 = affine.apply #map2(%i3, %i4, %i3, %i4)
          %13 = affine.apply #map3(%i3, %i4, %i3, %i4)
          store %11, %0[%12, %13] : memref<1x1xf32>
          %14 = load %arg0[%i3, %i4] : memref<2500x2500xf32>
          %15 = load %2[%i3, %i4] : memref<2500x2500xf32>
          %16 = addf %15, %14 : f32
          %17 = affine.apply #map2(%i3, %i4, %i3, %i4)
          %18 = affine.apply #map3(%i3, %i4, %i3, %i4)
          store %16, %1[%17, %18] : memref<1x1xf32>
          %19 = affine.apply #map2(%i3, %i4, %i3, %i4)
          %20 = affine.apply #map3(%i3, %i4, %i3, %i4)
          %21 = load %0[%19, %20] : memref<1x1xf32>
          %22 = affine.apply #map2(%i3, %i4, %i3, %i4)
          %23 = affine.apply #map3(%i3, %i4, %i3, %i4)
          %24 = load %1[%22, %23] : memref<1x1xf32>
          %25 = addf %24, %21 : f32
          store %25, %8[%i3, %i4] : memref<2500x2500xf32>
        }
      }
    
    * Rename MLIR_LLVM_OPTIONS to NGRAPH_MLIR_OPTIONS
    
    Something like this works now:
    NGRAPH_MLIR_OPTIONS="--enable-affine-loop-fusion=false"
    
    * Disable loop fusion by default and fix typo
    aedd8c2e
Name
Last commit
Last update
.ci Loading commit data...
cmake Loading commit data...
contrib/docker Loading commit data...
doc Loading commit data...
licenses Loading commit data...
maint Loading commit data...
python Loading commit data...
src Loading commit data...
test Loading commit data...
.clang-format Loading commit data...
.gitattributes Loading commit data...
.gitignore Loading commit data...
.travis.yml Loading commit data...
ABOUT.md Loading commit data...
CMakeLists.txt Loading commit data...
CODEOWNERS Loading commit data...
CONTRIB.md Loading commit data...
LICENSE Loading commit data...
README.md Loading commit data...
VERSION.in Loading commit data...
changes.md Loading commit data...
ecosystem-overview.md Loading commit data...