polymer
polymer copied to clipboard
Properly handle loop with arbitrary bounds
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
}
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.