1// RUN: mlir-opt -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s 2 3// CHECK: module attributes {gpu.container_module} 4 5// CHECK-LABEL: func @launch() 6func @launch() { 7 // CHECK: %[[ARG0:.*]] = "op"() : () -> f32 8 %0 = "op"() : () -> (f32) 9 // CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1> 10 %1 = "op"() : () -> (memref<?xf32, 1>) 11 // CHECK: %[[GDIMX:.*]] = constant 8 12 %gDimX = constant 8 : index 13 // CHECK: %[[GDIMY:.*]] = constant 12 14 %gDimY = constant 12 : index 15 // CHECK: %[[GDIMZ:.*]] = constant 16 16 %gDimZ = constant 16 : index 17 // CHECK: %[[BDIMX:.*]] = constant 20 18 %bDimX = constant 20 : index 19 // CHECK: %[[BDIMY:.*]] = constant 24 20 %bDimY = constant 24 : index 21 // CHECK: %[[BDIMZ:.*]] = constant 28 22 %bDimZ = constant 28 : index 23 24 // CHECK: "gpu.launch_func"(%[[GDIMX]], %[[GDIMY]], %[[GDIMZ]], %[[BDIMX]], %[[BDIMY]], %[[BDIMZ]], %[[ARG0]], %[[ARG1]]) {kernel = "launch_kernel", kernel_module = @launch_kernel} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> () 25 // CHECK-NOT: gpu.launch blocks 26 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY, 27 %grid_z = %gDimZ) 28 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY, 29 %block_z = %bDimZ) 30 args(%arg0 = %0, %arg1 = %1) : f32, memref<?xf32, 1> { 31 "use"(%arg0): (f32) -> () 32 "some_op"(%bx, %block_x) : (index, index) -> () 33 %42 = load %arg1[%tx] : memref<?xf32, 1> 34 gpu.return 35 } 36 return 37} 38 39 40// CHECK-LABEL: module @launch_kernel 41// CHECK-NEXT: gpu.func @launch_kernel 42// CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>) 43// CHECK-NEXT: %[[BID:.*]] = "gpu.block_id"() {dimension = "x"} : () -> index 44// CHECK-NEXT: = "gpu.block_id"() {dimension = "y"} : () -> index 45// CHECK-NEXT: = "gpu.block_id"() {dimension = "z"} : () -> index 46// CHECK-NEXT: %[[TID:.*]] = "gpu.thread_id"() {dimension = "x"} : () -> index 47// CHECK-NEXT: = "gpu.thread_id"() {dimension = "y"} : () -> index 48// CHECK-NEXT: = "gpu.thread_id"() {dimension = "z"} : () -> index 49// CHECK-NEXT: = "gpu.grid_dim"() {dimension = "x"} : () -> index 50// CHECK-NEXT: = "gpu.grid_dim"() {dimension = "y"} : () -> index 51// CHECK-NEXT: = "gpu.grid_dim"() {dimension = "z"} : () -> index 52// CHECK-NEXT: %[[BDIM:.*]] = "gpu.block_dim"() {dimension = "x"} : () -> index 53// CHECK-NEXT: = "gpu.block_dim"() {dimension = "y"} : () -> index 54// CHECK-NEXT: = "gpu.block_dim"() {dimension = "z"} : () -> index 55// CHECK-NEXT: "use"(%[[KERNEL_ARG0]]) : (f32) -> () 56// CHECK-NEXT: "some_op"(%[[BID]], %[[BDIM]]) : (index, index) -> () 57// CHECK-NEXT: = load %[[KERNEL_ARG1]][%[[TID]]] : memref<?xf32, 1> 58 59// ----- 60 61// CHECK: module attributes {gpu.container_module} 62 63func @multiple_launches() { 64 // CHECK: %[[CST:.*]] = constant 8 : index 65 %cst = constant 8 : index 66 // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = "multiple_launches_kernel", kernel_module = @multiple_launches_kernel} : (index, index, index, index, index, index) -> () 67 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst, 68 %grid_z = %cst) 69 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst, 70 %block_z = %cst) { 71 gpu.return 72 } 73 // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = "multiple_launches_kernel", kernel_module = @multiple_launches_kernel_0} : (index, index, index, index, index, index) -> () 74 gpu.launch blocks(%bx2, %by2, %bz2) in (%grid_x2 = %cst, %grid_y2 = %cst, 75 %grid_z2 = %cst) 76 threads(%tx2, %ty2, %tz2) in (%block_x2 = %cst, %block_y2 = %cst, 77 %block_z2 = %cst) { 78 gpu.return 79 } 80 return 81} 82 83// CHECK: module @multiple_launches_kernel 84// CHECK: func @multiple_launches_kernel 85// CHECK: module @multiple_launches_kernel_0 86// CHECK: func @multiple_launches_kernel 87 88// ----- 89 90func @extra_constants(%arg0 : memref<?xf32>) { 91 // CHECK: %[[CST:.*]] = constant 8 : index 92 %cst = constant 8 : index 93 %cst2 = constant 2 : index 94 %cst3 = dim %arg0, 0 : memref<?xf32> 95 // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %{{.*}}) {kernel = "extra_constants_kernel", kernel_module = @extra_constants_kernel} : (index, index, index, index, index, index, memref<?xf32>) -> () 96 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst, 97 %grid_z = %cst) 98 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst, 99 %block_z = %cst) 100 args(%kernel_arg0 = %cst2, %kernel_arg1 = %arg0, %kernel_arg2 = %cst3) : index, memref<?xf32>, index { 101 "use"(%kernel_arg0, %kernel_arg1, %kernel_arg2) : (index, memref<?xf32>, index) -> () 102 gpu.return 103 } 104 return 105} 106 107// CHECK-LABEL: func @extra_constants_kernel(%{{.*}}: memref<?xf32>) 108// CHECK: constant 109// CHECK: constant 110 111// ----- 112 113llvm.mlir.global internal @global(42 : i64) : !llvm.i64 114 115func @function_call(%arg0 : memref<?xf32>) { 116 %cst = constant 8 : index 117 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst, 118 %grid_z = %cst) 119 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst, 120 %block_z = %cst) { 121 call @device_function() : () -> () 122 call @device_function() : () -> () 123 %0 = llvm.mlir.addressof @global : !llvm<"i64*"> 124 gpu.return 125 } 126 return 127} 128 129func @device_function() { 130 call @recursive_device_function() : () -> () 131 gpu.return 132} 133 134func @recursive_device_function() { 135 call @recursive_device_function() : () -> () 136 gpu.return 137} 138 139// CHECK: module @function_call_kernel attributes {gpu.kernel_module} { 140// CHECK: gpu.func @function_call_kernel() 141// CHECK: call @device_function() : () -> () 142// CHECK: call @device_function() : () -> () 143// CHECK: llvm.mlir.addressof @global : !llvm<"i64*"> 144// 145// CHECK: llvm.mlir.global internal @global(42 : i64) : !llvm.i64 146// 147// CHECK: func @device_function() 148// CHECK: func @recursive_device_function() 149// CHECK-NOT: func @device_function 150