original_pw_conv_lower_to_loop - Joejiong/buddy-mlir GitHub Wiki

> /Workspace/buddy-mlir/llvm/build/bin/mlir-opt PointwiseConv2DNhwcHwcf.mlir -convert-linalg-to-loops                
#map = affine_map<(d0, d1) -> (d0 + d1)>
module  {
  func @pointwise_conv_2d_nhwc_hwcf(%arg0: memref<?x?x?x?xf32>, %arg1: memref<1x1x?x?xf32>, %arg2: memref<?x?x?x?xf32>) {
    %c0 = arith.constant 0 : index
    %c3 = arith.constant 3 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %0 = memref.dim %arg0, %c0 : memref<?x?x?x?xf32>
    %1 = memref.dim %arg0, %c3 : memref<?x?x?x?xf32>
    %2 = memref.dim %arg1, %c3 : memref<1x1x?x?xf32>
    %3 = memref.dim %arg2, %c1 : memref<?x?x?x?xf32>
    %4 = memref.dim %arg2, %c2 : memref<?x?x?x?xf32>
    scf.for %arg3 = %c0 to %0 step %c1 {
      scf.for %arg4 = %c0 to %3 step %c1 {
        scf.for %arg5 = %c0 to %4 step %c1 {
          scf.for %arg6 = %c0 to %2 step %c1 {
            scf.for %arg7 = %c0 to %c1 step %c1 {
              scf.for %arg8 = %c0 to %c1 step %c1 {
                scf.for %arg9 = %c0 to %1 step %c1 {
                  %5 = affine.apply #map(%arg4, %arg7)
                  %6 = affine.apply #map(%arg5, %arg8)
                  %7 = memref.load %arg0[%arg3, %5, %6, %arg9] : memref<?x?x?x?xf32>
                  %8 = memref.load %arg1[%arg7, %arg8, %arg9, %arg6] : memref<1x1x?x?xf32>
                  %9 = memref.load %arg2[%arg3, %arg4, %arg5, %arg6] : memref<?x?x?x?xf32>
                  %10 = arith.mulf %7, %8 : f32
                  %11 = arith.addf %9, %10 : f32
                  memref.store %11, %arg2[%arg3, %arg4, %arg5, %arg6] : memref<?x?x?x?xf32>
                }
              }
            }
          }
        }
      }
    }
    return
  }
}

> /Workspace/buddy-mlir/llvm/build/bin/mlir-opt PointwiseConv2DNhwcHwcf.mlir --convert-linalg-to-affine-loops        
#map = affine_map<(d0, d1) -> (d0 + d1)>
module  {
  func @pointwise_conv_2d_nhwc_hwcf(%arg0: memref<?x?x?x?xf32>, %arg1: memref<1x1x?x?xf32>, %arg2: memref<?x?x?x?xf32>) {
    %c0 = arith.constant 0 : index
    %c3 = arith.constant 3 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %0 = memref.dim %arg0, %c0 : memref<?x?x?x?xf32>
    %1 = memref.dim %arg0, %c3 : memref<?x?x?x?xf32>
    %2 = memref.dim %arg1, %c3 : memref<1x1x?x?xf32>
    %3 = memref.dim %arg2, %c1 : memref<?x?x?x?xf32>
    %4 = memref.dim %arg2, %c2 : memref<?x?x?x?xf32>
    affine.for %arg3 = 0 to %0 {
      affine.for %arg4 = 0 to %3 {
        affine.for %arg5 = 0 to %4 {
          affine.for %arg6 = 0 to %2 {
            affine.for %arg7 = 0 to 1 {
              affine.for %arg8 = 0 to 1 {
                affine.for %arg9 = 0 to %1 {
                  %5 = affine.apply #map(%arg4, %arg7)
                  %6 = affine.apply #map(%arg5, %arg8)
                  %7 = affine.load %arg0[%arg3, %5, %6, %arg9] : memref<?x?x?x?xf32>
                  %8 = affine.load %arg1[%arg7, %arg8, %arg9, %arg6] : memref<1x1x?x?xf32>
                  %9 = affine.load %arg2[%arg3, %arg4, %arg5, %arg6] : memref<?x?x?x?xf32>
                  %10 = arith.mulf %7, %8 : f32
                  %11 = arith.addf %9, %10 : f32
                  affine.store %11, %arg2[%arg3, %arg4, %arg5, %arg6] : memref<?x?x?x?xf32>
                }
              }
            }
          }
        }
      }
    }
    return
  }
}