101 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect -gpu-launch-sink-index-computations -split-input-file -verify-diagnostics %s | FileCheck %s2 3 4// CHECK-LABEL: @extra_constants5// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>6func.func @extra_constants(%arg0: memref<?xf32>) {7 %cst = arith.constant 8 : index8 %cst2 = arith.constant 2 : index9 %c0 = arith.constant 0 : index10 %cst3 = memref.dim %arg0, %c0 : memref<?xf32>11 // CHECK: gpu.launch blocks12 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,13 %grid_z = %cst)14 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,15 %block_z = %cst) {16 // CHECK-NOT: arith.constant 817 // CHECK: %[[CST2:.*]] = arith.constant 218 // CHECK-NEXT: %[[CST0:.*]] = arith.constant 019 // CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[ARG0]], %[[CST0]]20 // CHECK-NEXT: "use"(%[[CST2]], %[[ARG0]], %[[DIM]]) : (index, memref<?xf32>, index) -> ()21 // CHECK-NEXT: gpu.terminator22 "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()23 gpu.terminator24 }25 return26}27 28// -----29 30// CHECK-LABEL: @extra_constants_not_inlined31// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>32func.func @extra_constants_not_inlined(%arg0: memref<?xf32>) {33 %cst = arith.constant 8 : index34 %cst2 = arith.constant 2 : index35 %c0 = arith.constant 0 : index36 // CHECK: %[[CST_X:.*]] = "secret_constant"()37 %cst3 = "secret_constant"() : () -> index38 // CHECK: gpu.launch blocks39 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,40 %grid_z = %cst)41 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,42 %block_z = %cst) {43 // CHECK-NOT: arith.constant 844 // CHECK-NOT: "secret_constant"()45 // CHECK: %[[CST2:.*]] = arith.constant 246 // CHECK-NEXT: "use"(%[[CST2]], %[[ARG0]], %[[CST_X]]) : (index, memref<?xf32>, index) -> ()47 // CHECK-NEXT: gpu.terminator48 "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()49 gpu.terminator50 }51 return52}53 54// -----55 56// CHECK-LABEL: @multiple_uses57func.func @multiple_uses(%arg0 : memref<?xf32>) {58 %c1 = arith.constant 1 : index59 %c2 = arith.constant 2 : index60 // CHECK: gpu.launch blocks61 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,62 %grid_z = %c1)63 threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,64 %block_z = %c1) {65 // CHECK: %[[C2:.*]] = arith.constant 266 // CHECK-NEXT: "use1"(%[[C2]], %[[C2]])67 // CHECK-NEXT: "use2"(%[[C2]])68 // CHECK-NEXT: gpu.terminator69 "use1"(%c2, %c2) : (index, index) -> ()70 "use2"(%c2) : (index) -> ()71 gpu.terminator72 }73 return74}75 76// -----77 78// CHECK-LABEL: @multiple_uses279func.func @multiple_uses2(%arg0 : memref<*xf32>) {80 %c1 = arith.constant 1 : index81 %c2 = arith.constant 2 : index82 %d = memref.dim %arg0, %c2 : memref<*xf32>83 // CHECK: gpu.launch blocks84 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,85 %grid_z = %c1)86 threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,87 %block_z = %c1) {88 // CHECK: %[[C2:.*]] = arith.constant 2 : index89 // CHECK: %[[D:.*]] = memref.dim %[[ARG:.*]], %[[C2]]90 // CHECK: "use1"(%[[D]])91 // CHECK: "use2"(%[[C2]], %[[C2]])92 // CHECK: "use3"(%[[ARG]])93 // CHECK: gpu.terminator94 "use1"(%d) : (index) -> ()95 "use2"(%c2, %c2) : (index, index) -> ()96 "use3"(%arg0) : (memref<*xf32>) -> ()97 gpu.terminator98 }99 return100}101