47 lines · plain
1// RUN: mlir-opt %s -scf-parallel-loop-specialization -split-input-file | FileCheck %s2 3#map0 = affine_map<()[s0, s1] -> (1024, s0 - s1)>4#map1 = affine_map<()[s0, s1] -> (64, s0 - s1)>5 6func.func @parallel_loop(%outer_i0: index, %outer_i1: index, %A: memref<?x?xf32>, %B: memref<?x?xf32>,7 %C: memref<?x?xf32>, %result: memref<?x?xf32>) {8 %c0 = arith.constant 0 : index9 %c1 = arith.constant 1 : index10 %d0 = memref.dim %A, %c0 : memref<?x?xf32>11 %d1 = memref.dim %A, %c1 : memref<?x?xf32>12 %b0 = affine.min #map0()[%d0, %outer_i0]13 %b1 = affine.min #map1()[%d1, %outer_i1]14 scf.parallel (%i0, %i1) = (%c0, %c0) to (%b0, %b1) step (%c1, %c1) {15 %B_elem = memref.load %B[%i0, %i1] : memref<?x?xf32>16 %C_elem = memref.load %C[%i0, %i1] : memref<?x?xf32>17 %sum_elem = arith.addf %B_elem, %C_elem : f3218 memref.store %sum_elem, %result[%i0, %i1] : memref<?x?xf32>19 }20 return21}22 23// CHECK-LABEL: func @parallel_loop(24// CHECK-SAME: [[VAL_0:%.*]]: index, [[VAL_1:%.*]]: index, [[VAL_2:%.*]]: memref<?x?xf32>, [[VAL_3:%.*]]: memref<?x?xf32>, [[VAL_4:%.*]]: memref<?x?xf32>, [[VAL_5:%.*]]: memref<?x?xf32>) {25// CHECK: [[VAL_6:%.*]] = arith.constant 0 : index26// CHECK: [[VAL_7:%.*]] = arith.constant 1 : index27// CHECK: [[VAL_8:%.*]] = memref.dim [[VAL_2]], [[VAL_6]] : memref<?x?xf32>28// CHECK: [[VAL_9:%.*]] = memref.dim [[VAL_2]], [[VAL_7]] : memref<?x?xf32>29// CHECK: [[VAL_10:%.*]] = affine.min #{{.*}}(){{\[}}[[VAL_8]], [[VAL_0]]]30// CHECK: [[VAL_11:%.*]] = affine.min #{{.*}}(){{\[}}[[VAL_9]], [[VAL_1]]]31// CHECK: [[VAL_12:%.*]] = arith.constant 1024 : index32// CHECK: [[VAL_13:%.*]] = arith.cmpi eq, [[VAL_10]], [[VAL_12]] : index33// CHECK: [[VAL_14:%.*]] = arith.constant 64 : index34// CHECK: [[VAL_15:%.*]] = arith.cmpi eq, [[VAL_11]], [[VAL_14]] : index35// CHECK: [[VAL_16:%.*]] = arith.andi [[VAL_13]], [[VAL_15]] : i136// CHECK: scf.if [[VAL_16]] {37// CHECK: scf.parallel ([[VAL_17:%.*]], [[VAL_18:%.*]]) = ([[VAL_6]], [[VAL_6]]) to ([[VAL_12]], [[VAL_14]]) step ([[VAL_7]], [[VAL_7]]) {38// CHECK: memref.store39// CHECK: }40// CHECK: } else {41// CHECK: scf.parallel ([[VAL_22:%.*]], [[VAL_23:%.*]]) = ([[VAL_6]], [[VAL_6]]) to ([[VAL_10]], [[VAL_11]]) step ([[VAL_7]], [[VAL_7]]) {42// CHECK: memref.store43// CHECK: }44// CHECK: }45// CHECK: return46// CHECK: }47