kumasento / polymer

Bridging polyhedral analysis tools to the MLIR framework
MIT License
99 stars 20 forks source link

Added extract-scop-stmt pass #60

Closed kumasento closed 3 years ago

kumasento commented 3 years ago

This PR adds a new pass -extract-scop-stmt, which will be used to discover scop statements in the input code and turn them into function calls. This will be handy when translating MLIR to/from OpenScop.

A use case:

func @load_multi_stores(%A: memref<?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %N = dim %B, %c0 : memref<?x?xf32> 
  %M = dim %B, %c1 : memref<?x?xf32> 

  affine.for %i = 0 to %N {
    affine.for %j = 0 to %M {
      %0 = affine.load %A[%i] : memref<?xf32>
      %1 = mulf %0, %0 : f32
      affine.store %1, %B[%i, %j] : memref<?x?xf32>
      %2 = addf %0, %1 : f32
      affine.store %2, %C[%i, %j] : memref<?x?xf32>
    }
  }

  return
}

Result:

  func @load_multi_stores(%arg0: memref<?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
    %c0 = constant 0 : index
    %c1 = constant 1 : index
    %0 = dim %arg1, %c0 : memref<?x?xf32>
    %1 = dim %arg1, %c1 : memref<?x?xf32>
    affine.for %arg3 = 0 to %0 {
      affine.for %arg4 = 0 to %1 {
        call @S0(%arg0, %arg3, %arg1, %arg4) : (memref<?xf32>, index, memref<?x?xf32>, index) -> ()
        call @S1(%arg0, %arg3, %arg2, %arg4) : (memref<?xf32>, index, memref<?x?xf32>, index) -> ()
      }
    }
    return
  }
  func @S0(%arg0: memref<?xf32>, %arg1: index, %arg2: memref<?x?xf32>, %arg3: index) {
    %0 = affine.load %arg0[%arg1] : memref<?xf32>
    %1 = mulf %0, %0 : f32
    affine.store %1, %arg2[%arg1, %arg3] : memref<?x?xf32>
    return
  }
  func @S1(%arg0: memref<?xf32>, %arg1: index, %arg2: memref<?x?xf32>, %arg3: index) {
    %0 = affine.load %arg0[%arg1] : memref<?xf32>
    %1 = mulf %0, %0 : f32
    %2 = addf %0, %1 : f32
    affine.store %2, %arg2[%arg1, %arg3] : memref<?x?xf32>
    return
  }