kumasento / polymer

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

Properly handle loop with arbitrary bounds #42

Open kumasento opened 3 years ago

kumasento commented 3 years ago

It would be pretty common in real-world test cases.

Test code -

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

  %M = dim %A, %c0 : memref<?x?xf32>
  %N = dim %B, %c0 : memref<?x?xf32>
  %K = dim %A, %c1 : memref<?x?xf32>

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

  return 
}

func @main() {
  %M = constant 64 : index
  %N = constant 64 : index
  %K = constant 64 : index

  %A = alloc(%M, %K) : memref<?x?xf32>
  %B = alloc(%N, %K) : memref<?x?xf32>
  %C = alloc(%M, %N) : memref<?x?xf32>

  call @matmul(%A, %B, %C)
   : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()

  return
}
kumasento commented 3 years ago

Will be merged from the impact branch pretty soon, but the key idea is: we will recreate those dim operations in the new piece of code.