1// RUN: mlir-opt -allow-unregistered-dialect -gpu-launch-sink-index-computations -split-input-file -verify-diagnostics %s | FileCheck %s 2 3 4// CHECK-LABEL: @extra_constants 5// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32> 6func.func @extra_constants(%arg0: memref<?xf32>) { 7 %cst = arith.constant 8 : index 8 %cst2 = arith.constant 2 : index 9 %c0 = arith.constant 0 : index 10 %cst3 = memref.dim %arg0, %c0 : memref<?xf32> 11 // CHECK: gpu.launch blocks 12 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 8 17 // CHECK: %[[CST2:.*]] = arith.constant 2 18 // CHECK-NEXT: %[[CST0:.*]] = arith.constant 0 19 // CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[ARG0]], %[[CST0]] 20 // CHECK-NEXT: "use"(%[[CST2]], %[[ARG0]], %[[DIM]]) : (index, memref<?xf32>, index) -> () 21 // CHECK-NEXT: gpu.terminator 22 "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> () 23 gpu.terminator 24 } 25 return 26} 27 28// ----- 29 30// CHECK-LABEL: @extra_constants_not_inlined 31// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32> 32func.func @extra_constants_not_inlined(%arg0: memref<?xf32>) { 33 %cst = arith.constant 8 : index 34 %cst2 = arith.constant 2 : index 35 %c0 = arith.constant 0 : index 36 // CHECK: %[[CST_X:.*]] = "secret_constant"() 37 %cst3 = "secret_constant"() : () -> index 38 // CHECK: gpu.launch blocks 39 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 8 44 // CHECK-NOT: "secret_constant"() 45 // CHECK: %[[CST2:.*]] = arith.constant 2 46 // CHECK-NEXT: "use"(%[[CST2]], %[[ARG0]], %[[CST_X]]) : (index, memref<?xf32>, index) -> () 47 // CHECK-NEXT: gpu.terminator 48 "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> () 49 gpu.terminator 50 } 51 return 52} 53 54// ----- 55 56// CHECK-LABEL: @multiple_uses 57func.func @multiple_uses(%arg0 : memref<?xf32>) { 58 %c1 = arith.constant 1 : index 59 %c2 = arith.constant 2 : index 60 // CHECK: gpu.launch blocks 61 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 2 66 // CHECK-NEXT: "use1"(%[[C2]], %[[C2]]) 67 // CHECK-NEXT: "use2"(%[[C2]]) 68 // CHECK-NEXT: gpu.terminator 69 "use1"(%c2, %c2) : (index, index) -> () 70 "use2"(%c2) : (index) -> () 71 gpu.terminator 72 } 73 return 74} 75 76// ----- 77 78// CHECK-LABEL: @multiple_uses2 79func.func @multiple_uses2(%arg0 : memref<*xf32>) { 80 %c1 = arith.constant 1 : index 81 %c2 = arith.constant 2 : index 82 %d = memref.dim %arg0, %c2 : memref<*xf32> 83 // CHECK: gpu.launch blocks 84 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 : index 89 // CHECK: %[[D:.*]] = memref.dim %[[ARG:.*]], %[[C2]] 90 // CHECK: "use1"(%[[D]]) 91 // CHECK: "use2"(%[[C2]], %[[C2]]) 92 // CHECK: "use3"(%[[ARG]]) 93 // CHECK: gpu.terminator 94 "use1"(%d) : (index) -> () 95 "use2"(%c2, %c2) : (index, index) -> () 96 "use3"(%arg0) : (memref<*xf32>) -> () 97 gpu.terminator 98 } 99 return 100} 101