202 lines · plain
1// RUN: mlir-opt %s --transform-interpreter | FileCheck %s2// RUN: mlir-opt %s --gpu-eliminate-barriers | FileCheck %s3 4module attributes {transform.with_named_sequence} {5 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {6 %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op7 transform.apply_patterns to %0 {8 transform.apply_patterns.gpu.eliminate_barriers9 } : !transform.any_op10 transform.yield11 }12}13 14// CHECK-LABEL: @read_read_write15func.func @read_read_write(%arg0: memref<?xf32>, %arg1: index) attributes {__parallel_region_boundary_for_test} {16 // CHECK: load17 %0 = memref.load %arg0[%arg1] : memref<?xf32>18 // The barrier between loads can be removed.19 // CHECK-NOT: barrier20 gpu.barrier21 // CHECK: load22 %1 = memref.load %arg0[%arg1] : memref<?xf32>23 %2 = arith.addf %0, %1 : f3224 // The barrier between load and store cannot be removed (unless we reason about accessed subsets).25 // CHECK: barrier26 gpu.barrier27 // CHECK: store28 memref.store %2, %arg0[%arg1] : memref<?xf32>29 return30}31 32// CHECK-LABEL: @write_read_read33func.func @write_read_read(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) -> f3234attributes {__parallel_region_boundary_for_test} {35 // CHECK: store36 memref.store %arg2, %arg0[%arg1] : memref<?xf32>37 // The barrier between load and store cannot be removed (unless we reason about accessed subsets).38 // CHECK: barrier39 gpu.barrier40 // CHECK: load41 %0 = memref.load %arg0[%arg1] : memref<?xf32>42 // CHECK-NOT: barrier43 gpu.barrier44 // CHECK: load45 %1 = memref.load %arg0[%arg1] : memref<?xf32>46 %2 = arith.addf %0, %1 : f3247 return %2 : f3248}49 50// CHECK-LABEL: @write_in_a_loop51func.func @write_in_a_loop(%arg0: memref<?xf32>, %arg1: f32) attributes {__parallel_region_boundary_for_test} {52 %c0 = arith.constant 0 : index53 %c42 = arith.constant 42 : index54 %c1 = arith.constant 1 : index55 scf.for %i = %c0 to %c42 step %c1 {56 memref.store %arg1, %arg0[%i] : memref<?xf32>57 // Cannot remove this barrier because it guards write-after-write between different iterations.58 // CHECK: barrier59 gpu.barrier60 }61 return62}63 64// CHECK-LABEL: @read_read_write_loop65func.func @read_read_write_loop(%arg0: memref<?xf32>, %arg1: f32) attributes {__parallel_region_boundary_for_test} {66 %c0 = arith.constant 0 : index67 %c42 = arith.constant 42 : index68 %c1 = arith.constant 1 : index69 scf.for %i = %c0 to %c42 step %c1 {70 // (Note that if subscript were different, this would have been a race with the store at the end of the loop).71 %0 = memref.load %arg0[%i] : memref<?xf32>72 // Guards read-after-write where the write happens on the previous iteration.73 // CHECK: barrier74 gpu.barrier75 %1 = memref.load %arg0[%i] : memref<?xf32>76 %2 = arith.addf %0, %1 : f3277 // Guards write-after-read.78 // CHECK: barrier79 gpu.barrier80 memref.store %2, %arg0[%i] : memref<?xf32>81 }82 return83}84 85// CHECK-LABEL: @read_read_write_loop_trailing_sync86func.func @read_read_write_loop_trailing_sync(%arg0: memref<?xf32>, %arg1: f32) attributes {__parallel_region_boundary_for_test} {87 %c0 = arith.constant 0 : index88 %c42 = arith.constant 42 : index89 %c1 = arith.constant 1 : index90 scf.for %i = %c0 to %c42 step %c1 {91 // CHECK: load92 %0 = memref.load %arg0[%i] : memref<?xf32>93 // This can be removed because it only guards a read-after-read.94 // CHECK-NOT: barrier95 gpu.barrier96 // CHECK: load97 %1 = memref.load %arg0[%i] : memref<?xf32>98 %2 = arith.addf %0, %1 : f3299 // CHECK: barrier100 gpu.barrier101 // CHECK: store102 memref.store %2, %arg0[%i] : memref<?xf32>103 // CHECK: barrier104 gpu.barrier105 }106 return107}108 109// CHECK-LABEL: @write_write_noalias110func.func @write_write_noalias(%arg0: index, %arg1: f32) -> (memref<42xf32>, memref<10xf32>)111attributes {__parallel_region_boundary_for_test} {112 %0 = memref.alloc() : memref<42xf32>113 %1 = memref.alloc() : memref<10xf32>114 // CHECK: store115 memref.store %arg1, %0[%arg0] : memref<42xf32>116 // This can be removed because we can prove two allocations don't alias.117 // CHECK-NOT: barrier118 gpu.barrier119 // CHECK: store120 memref.store %arg1, %1[%arg0] : memref<10xf32>121 return %0, %1 : memref<42xf32>, memref<10xf32>122}123 124// CHECK-LABEL: @write_write_alloc_arg_noalias125func.func @write_write_alloc_arg_noalias(%arg0: index, %arg1: f32, %arg2: memref<?xf32>) -> (memref<42xf32>)126attributes {__parallel_region_boundary_for_test} {127 %0 = memref.alloc() : memref<42xf32>128 // CHECK: store129 memref.store %arg1, %0[%arg0] : memref<42xf32>130 // This can be removed because we can prove local allocation doesn't alias with a function argument.131 // CHECK-NOT: barrier132 gpu.barrier133 // CHECK: store134 memref.store %arg1, %arg2[%arg0] : memref<?xf32>135 return %0 : memref<42xf32>136}137 138// CHECK-LABEL: @repeated_barrier139func.func @repeated_barrier(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) -> f32140attributes {__parallel_region_boundary_for_test} {141 %0 = memref.load %arg0[%arg1] : memref<?xf32>142 // CHECK: gpu.barrier143 gpu.barrier144 // CHECK-NOT: gpu.barrier145 gpu.barrier146 memref.store %arg2, %arg0[%arg1] : memref<?xf32>147 return %0 : f32148}149 150// CHECK-LABEL: @symmetric_stop151func.func @symmetric_stop(%val: f32) -> (f32, f32, f32, f32, f32)152attributes {__parallel_region_boundary_for_test} {153 // CHECK: %[[A:.+]] = memref.alloc154 // CHECK: %[[B:.+]] = memref.alloc155 // CHECK: %[[C:.+]] = memref.alloc156 %A = memref.alloc() : memref<f32>157 %B = memref.alloc() : memref<f32>158 %C = memref.alloc() : memref<f32>159 // CHECK: memref.store %{{.*}}, %[[A]]160 memref.store %val, %A[] : memref<f32>161 // CHECK: gpu.barrier162 gpu.barrier163 // CHECK: memref.load %[[A]]164 %0 = memref.load %A[] : memref<f32>165 // CHECK: memref.store %{{.*}}, %[[B]]166 memref.store %val, %B[] : memref<f32>167 // This barrier is eliminated because the surrounding barriers are sufficient168 // to guard write/read on all memrefs.169 // CHECK-NOT: gpu.barrier170 gpu.barrier171 // CHECK: memref.load %[[A]]172 %1 = memref.load %A[] : memref<f32>173 // CHECK: memref.store %{{.*}} %[[C]]174 memref.store %val, %C[] : memref<f32>175 // CHECK: gpu.barrier176 gpu.barrier177 // CHECK: memref.load %[[A]]178 // CHECK: memref.load %[[B]]179 // CHECK: memref.load %[[C]]180 %2 = memref.load %A[] : memref<f32>181 %3 = memref.load %B[] : memref<f32>182 %4 = memref.load %C[] : memref<f32>183 return %0, %1, %2, %3, %4 : f32, f32, f32, f32, f32184}185 186// CHECK-LABEL: @nested_loop_barrier_only187func.func @nested_loop_barrier_only() attributes {__parallel_region_boundary_for_test} {188 %c0 = arith.constant 0 : index189 %c42 = arith.constant 42 : index190 %c1 = arith.constant 1 : index191 // Note: the barrier can be removed and as consequence the loops get folded192 // by the greedy rewriter.193 // CHECK-NOT: scf.for194 // CHECK-NOT: gpu.barrier195 scf.for %j = %c0 to %c42 step %c1 {196 scf.for %i = %c0 to %c42 step %c1 {197 gpu.barrier198 }199 }200 return201}202