Dialect/SCF/parallel-loop-fusion.mlir

27 func @fuse_two(%A: memref<2x2xf32>, %B: memref<2x2xf32>,
28                     %C: memref<2x2xf32>, %result: memref<2x2xf32>) {
32   %sum = memref.alloc()  : memref<2x2xf32>
34     %B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
35     %C_elem = memref.load %C[%i, %j] : memref<2x2xf32>
37     memref.store %sum_elem, %sum[%i, %j] : memref<2x2xf32>
41     %sum_elem = memref.load %sum[%i, %j] : memref<2x2xf32>
42     %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
44     memref.store %product_elem, %result[%i, %j] : memref<2x2xf32>
47   memref.dealloc %sum : memref<2x2xf32>
56 // CHECK:      [[SUM:%.*]] = memref.alloc()
59 // CHECK:        [[B_ELEM:%.*]] = memref.load [[B]]{{\[}}[[I]], [[J]]]
60 // CHECK:        [[C_ELEM:%.*]] = memref.load [[C]]{{\[}}[[I]], [[J]]]
62 // CHECK:        memref.store [[SUM_ELEM]], [[SUM]]{{\[}}[[I]], [[J]]]
63 // CHECK:        [[SUM_ELEM_:%.*]] = memref.load [[SUM]]{{\[}}[[I]], [[J]]]
64 // CHECK:        [[A_ELEM:%.*]] = memref.load [[A]]{{\[}}[[I]], [[J]]]
66 // CHECK:        memref.store [[PRODUCT_ELEM]], [[RESULT]]{{\[}}[[I]], [[J]]]
69 // CHECK:      memref.dealloc [[SUM]]
73 func @fuse_three(%lhs: memref<100x10xf32>, %rhs: memref<100xf32>,
74                       %result: memref<100x10xf32>) {
79   %broadcast_rhs = memref.alloc() : memref<100x10xf32>
80   %diff = memref.alloc() : memref<100x10xf32>
82     %rhs_elem = memref.load %rhs[%i] : memref<100xf32>
83     memref.store %rhs_elem, %broadcast_rhs[%i, %j] : memref<100x10xf32>
87     %lhs_elem = memref.load %lhs[%i, %j] : memref<100x10xf32>
88     %broadcast_rhs_elem = memref.load %broadcast_rhs[%i, %j] : memref<100x10xf32>
90     memref.store %diff_elem, %diff[%i, %j] : memref<100x10xf32>
94     %diff_elem = memref.load %diff[%i, %j] : memref<100x10xf32>
96     memref.store %exp_elem, %result[%i, %j] : memref<100x10xf32>
99   memref.dealloc %broadcast_rhs : memref<100x10xf32>
100   memref.dealloc %diff : memref<100x10xf32>
104 // CHECK-SAME: ([[LHS:%.*]]: memref<100x10xf32>, [[RHS:%.*]]: memref<100xf32>,
105 // CHECK-SAME:  [[RESULT:%.*]]: memref<100x10xf32>) {
110 // CHECK:      [[BROADCAST_RHS:%.*]] = memref.alloc()
111 // CHECK:      [[DIFF:%.*]] = memref.alloc()
114 // CHECK:        [[RHS_ELEM:%.*]] = memref.load [[RHS]]{{\[}}[[I]]]
115 // CHECK:        memref.store [[RHS_ELEM]], [[BROADCAST_RHS]]{{\[}}[[I]], [[J]]]
116 // CHECK:        [[LHS_ELEM:%.*]] = memref.load [[LHS]]{{\[}}[[I]], [[J]]]
117 // CHECK:        [[BROADCAST_RHS_ELEM:%.*]] = memref.load [[BROADCAST_RHS]]
119 // CHECK:        memref.store [[DIFF_ELEM]], [[DIFF]]{{\[}}[[I]], [[J]]]
120 // CHECK:        [[DIFF_ELEM_:%.*]] = memref.load [[DIFF]]{{\[}}[[I]], [[J]]]
122 // CHECK:        memref.store [[EXP_ELEM]], [[RESULT]]{{\[}}[[I]], [[J]]]
125 // CHECK:      memref.dealloc [[BROADCAST_RHS]]
126 // CHECK:      memref.dealloc [[DIFF]]
199   %buffer  = memref.alloc() : memref<2x2xf32>
231     %A: memref<2x2xf32>, %B: memref<2x2xf32>,
232     %C: memref<2x2xf32>, %result: memref<2x2xf32>) {
236   %common_buf = memref.alloc() : memref<2x2xf32>
238     %B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
239     %C_elem = memref.load %C[%i, %j] : memref<2x2xf32>
241     memref.store %sum_elem, %common_buf[%i, %j] : memref<2x2xf32>
246     %sum_elem = memref.load %common_buf[%k, %j] : memref<2x2xf32>
247     %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
249     memref.store %product_elem, %result[%i, %j] : memref<2x2xf32>
252   memref.dealloc %common_buf : memref<2x2xf32>
262     %A: memref<2x2xf32>, %B: memref<2x2xf32>, %common_buf: memref<2x2xf32>) {
266   %sum = memref.alloc() : memref<2x2xf32>
268     %B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
269     %C_elem = memref.load %common_buf[%i, %j] : memref<2x2xf32>
271     memref.store %sum_elem, %sum[%i, %j] : memref<2x2xf32>
276     %sum_elem = memref.load %sum[%k, %j] : memref<2x2xf32>
277     %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
279     memref.store %product_elem, %common_buf[%j, %i] : memref<2x2xf32>
282   memref.dealloc %sum : memref<2x2xf32>
295   %buffer  = memref.alloc() : memref<2x2xf32>
300     %A = memref.subview %buffer[%c0, %c0][%c2, %c2][%c1, %c1]
301       : memref<2x2xf32> to memref<?x?xf32, offset: ?, strides:[?, ?]>
302     %A_elem = memref.load %A[%i, %j] : memref<?x?xf32, offset: ?, strides:[?, ?]>
313 func @nested_fuse(%A: memref<2x2xf32>, %B: memref<2x2xf32>,
314                     %C: memref<2x2xf32>, %result: memref<2x2xf32>) {
318   %sum = memref.alloc()  : memref<2x2xf32>
321       %B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
322       %C_elem = memref.load %C[%i, %j] : memref<2x2xf32>
324       memref.store %sum_elem, %sum[%i, %j] : memref<2x2xf32>
328       %sum_elem = memref.load %sum[%i, %j] : memref<2x2xf32>
329       %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
331       memref.store %product_elem, %result[%i, %j] : memref<2x2xf32>
335   memref.dealloc %sum : memref<2x2xf32>
344 // CHECK:      [[SUM:%.*]] = memref.alloc()
348 // CHECK:          [[B_ELEM:%.*]] = memref.load [[B]]{{\[}}[[I]], [[J]]]
349 // CHECK:          [[C_ELEM:%.*]] = memref.load [[C]]{{\[}}[[I]], [[J]]]
351 // CHECK:          memref.store [[SUM_ELEM]], [[SUM]]{{\[}}[[I]], [[J]]]
352 // CHECK:          [[SUM_ELEM_:%.*]] = memref.load [[SUM]]{{\[}}[[I]], [[J]]]
353 // CHECK:          [[A_ELEM:%.*]] = memref.load [[A]]{{\[}}[[I]], [[J]]]
355 // CHECK:          memref.store [[PRODUCT_ELEM]], [[RESULT]]{{\[}}[[I]], [[J]]]
359 // CHECK:      memref.dealloc [[SUM]]