llvm-project
2ecf6088 - [mlir]Fix compose subview (#80551)

Commit
2 years ago
[mlir]Fix compose subview (#80551) I found a bug in `test-compose-subview`,You can see the example I gave. ``` #map = affine_map<() -> ()> module { func.func private @fun(%arg0: memref<10x10xf32>, %arg1: memref<5x5xf32>) -> memref<5x5xf32> { %c0 = arith.constant 0 : index %c5 = arith.constant 5 : index %c1 = arith.constant 1 : index %subview = memref.subview %arg0[0, 0] [5, 5] [1, 1] : memref<10x10xf32> to memref<5x5xf32, strided<[10, 1]>> %alloc = memref.alloc() : memref<5x5xf32> scf.for %arg2 = %c0 to %c5 step %c1 { scf.for %arg3 = %c0 to %c5 step %c1 { %subview_0 = memref.subview %subview[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32, strided<[10, 1]>> to memref<f32, strided<[], offset: ?>> %subview_1 = memref.subview %arg1[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32> to memref<f32, strided<[], offset: ?>> %alloc_2 = memref.alloc() : memref<f32> linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} ins(%subview_0, %subview_1 : memref<f32, strided<[], offset: ?>>, memref<f32, strided<[], offset: ?>>) outs(%alloc_2 : memref<f32>) { ^bb0(%in: f32, %in_4: f32, %out: f32): %0 = arith.addf %in, %in_4 : f32 linalg.yield %0 : f32 } %subview_3 = memref.subview %alloc[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32> to memref<f32, strided<[], offset: ?>> memref.copy %alloc_2, %subview_3 : memref<f32> to memref<f32, strided<[], offset: ?>> } } return %alloc : memref<5x5xf32> } func.func @test(%arg0: memref<10x10xf32>, %arg1: memref<5x5xf32>) -> memref<5x5xf32> { %0 = call @fun(%arg0, %arg1) : (memref<10x10xf32>, memref<5x5xf32>) -> memref<5x5xf32> return %0 : memref<5x5xf32> } } ``` When I run `mlir-opt test.mlir ---test-compose-subview`. ``` test.mlir:14:9: error: 'linalg.generic' op expected operand rank (2) to match the result rank of indexing_map #0 (0) linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} ins(%subview_0, %subview_1 : memref<f32, strided<[], offset: ?>>, memref<f32, strided<[], offset: ?>>) outs(%alloc_2 : memref<f32>) { ^ test1.mlir:14:9: note: see current operation: "linalg.generic"(%4, %5, %6) <{indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = [], operandSegmentSizes = array<i32: 2, 1>}> ({ ^bb0(%arg4: f32, %arg5: f32, %arg6: f32): %8 = "arith.addf"(%arg4, %arg5) <{fastmath = #arith.fastmath<none>}> : (f32, f32) -> f32 "linalg.yield"(%8) : (f32) -> () }) : (memref<1x1xf32, strided<[10, 1], offset: ?>>, memref<f32, strided<[], offset: ?>>, memref<f32>) -> () ``` This PR fixes that.In the meantime I've extended this PR to handle cases where stride is greater than 1. ``` func.func private @Unknown0(%arg0: memref<10x10xf32>, %arg1: memref<5x5xf32>) -> memref<5x5xf32> { %c0 = arith.constant 0 : index %c5 = arith.constant 5 : index %c1 = arith.constant 1 : index %subview = memref.subview %arg0[0, 0] [5, 5] [2, 2] : memref<10x10xf32> to memref<5x5xf32, strided<[20, 2]>> %alloc = memref.alloc() : memref<5x5xf32> scf.for %arg2 = %c0 to %c5 step %c1 { scf.for %arg3 = %c0 to %c5 step %c1 { %subview_0 = memref.subview %subview[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32, strided<[20, 2]>> to memref<f32, strided<[], offset: ?>> %subview_1 = memref.subview %arg1[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32> to memref<f32, strided<[], offset: ?>> %alloc_2 = memref.alloc() : memref<f32> linalg.generic {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} ins(%subview_0, %subview_1 : memref<f32, strided<[], offset: ?>>, memref<f32, strided<[], offset: ?>>) outs(%alloc_2 : memref<f32>) { ^bb0(%in: f32, %in_4: f32, %out: f32): %0 = arith.addf %in, %in_4 : f32 linalg.yield %0 : f32 } %subview_3 = memref.subview %alloc[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32> to memref<f32, strided<[], offset: ?>> memref.copy %alloc_2, %subview_3 : memref<f32> to memref<f32, strided<[], offset: ?>> } } return %alloc : memref<5x5xf32> } $ mlir-opt test.mlir -test-compose-subview #map = affine_map<()[s0] -> (s0 * 2)> #map1 = affine_map<() -> ()> module { func.func private @Unknown0(%arg0: memref<10x10xf32>, %arg1: memref<5x5xf32>) -> memref<5x5xf32> { %c0 = arith.constant 0 : index %c5 = arith.constant 5 : index %c1 = arith.constant 1 : index %alloc = memref.alloc() : memref<5x5xf32> scf.for %arg2 = %c0 to %c5 step %c1 { scf.for %arg3 = %c0 to %c5 step %c1 { %0 = affine.apply #map()[%arg2] %1 = affine.apply #map()[%arg3] %subview = memref.subview %arg0[%0, %1] [1, 1] [2, 2] : memref<10x10xf32> to memref<f32, strided<[], offset: ?>> %subview_0 = memref.subview %arg1[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32> to memref<f32, strided<[], offset: ?>> %alloc_1 = memref.alloc() : memref<f32> linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = []} ins(%subview, %subview_0 : memref<f32, strided<[], offset: ?>>, memref<f32, strided<[], offset: ?>>) outs(%alloc_1 : memref<f32>) { ^bb0(%in: f32, %in_3: f32, %out: f32): %2 = arith.addf %in, %in_3 : f32 linalg.yield %2 : f32 } %subview_2 = memref.subview %alloc[%arg2, %arg3] [1, 1] [1, 1] : memref<5x5xf32> to memref<f32, strided<[], offset: ?>> memref.copy %alloc_1, %subview_2 : memref<f32> to memref<f32, strided<[], offset: ?>> } } return %alloc : memref<5x5xf32> } } ```
Parents
Loading