mlir/test/Conversion/GPUToSPIRV/builtins.mlir

   1 // RUN: mlir-opt -split-input-file -convert-gpu-to-spirv %s -o - | FileCheck %s
   2
   3 module attributes {gpu.container_module} {
   4   func @builtin() {
   5     %c0 = arith.constant 1 : index
   6     gpu.launch_func @kernels::@builtin_workgroup_id_x
   7         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
   8     return
   9   }
  10
  11   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
  12   // CHECK: spv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
  13   gpu.module @kernels {
  14     gpu.func @builtin_workgroup_id_x() kernel
  15       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
  16       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPID]]
  17       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
  18       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
  19       %0 = "gpu.block_id"() {dimension = "x"} : () -> index
  20       gpu.return
  21     }
  22   }
  23 }
  24
  25 // -----
  26
  27 module attributes {gpu.container_module} {
  28   func @builtin() {
  29     %c0 = arith.constant 1 : index
  30     %c256 = arith.constant 256 : i32
  31     gpu.launch_func @kernels::@builtin_workgroup_id_y
  32         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
  33         dynamic_shared_memory_size %c256
  34     return
  35   }
  36
  37   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
  38   // CHECK: spv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
  39   gpu.module @kernels {
  40     gpu.func @builtin_workgroup_id_y() kernel
  41       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
  42       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPID]]
  43       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
  44       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}}
  45       %0 = "gpu.block_id"() {dimension = "y"} : () -> index
  46       gpu.return
  47     }
  48   }
  49 }
  50
  51 // -----
  52
  53 module attributes {gpu.container_module} {
  54   func @builtin() {
  55     %c0 = arith.constant 1 : index
  56     gpu.launch_func @kernels::@builtin_workgroup_id_z
  57         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
  58     return
  59   }
  60
  61   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
  62   // CHECK: spv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
  63   gpu.module @kernels {
  64     gpu.func @builtin_workgroup_id_z() kernel
  65       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
  66       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPID]]
  67       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
  68       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}2 : i32{{\]}}
  69       %0 = "gpu.block_id"() {dimension = "z"} : () -> index
  70       gpu.return
  71     }
  72   }
  73 }
  74
  75 // -----
  76
  77 module attributes {gpu.container_module} {
  78   func @builtin() {
  79     %c0 = arith.constant 1 : index
  80     gpu.launch_func @kernels::@builtin_workgroup_size_x
  81         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
  82     return
  83   }
  84
  85   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
  86   gpu.module @kernels {
  87     gpu.func @builtin_workgroup_size_x() kernel
  88       attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]>: vector<3xi32>}} {
  89       // The constant value is obtained from the spv.entry_point_abi.
  90       // Note that this ignores the workgroup size specification in gpu.launch.
  91       // We may want to define gpu.workgroup_size and convert it to the entry
  92       // point ABI we want here.
  93       // CHECK: spv.Constant 32 : i32
  94       %0 = "gpu.block_dim"() {dimension = "x"} : () -> index
  95       gpu.return
  96     }
  97   }
  98 }
  99
 100 // -----
 101
 102 module attributes {gpu.container_module} {
 103   func @builtin() {
 104     %c0 = arith.constant 1 : index
 105     gpu.launch_func @kernels::@builtin_workgroup_size_y
 106         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
 107     return
 108   }
 109
 110   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
 111   gpu.module @kernels {
 112     gpu.func @builtin_workgroup_size_y() kernel
 113       attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
 114       // The constant value is obtained from the spv.entry_point_abi.
 115       // CHECK: spv.Constant 4 : i32
 116       %0 = "gpu.block_dim"() {dimension = "y"} : () -> index
 117       gpu.return
 118     }
 119   }
 120 }
 121
 122 // -----
 123
 124 module attributes {gpu.container_module} {
 125   func @builtin() {
 126     %c0 = arith.constant 1 : index
 127     gpu.launch_func @kernels::@builtin_workgroup_size_z
 128         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
 129     return
 130   }
 131
 132   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
 133   gpu.module @kernels {
 134     gpu.func @builtin_workgroup_size_z() kernel
 135       attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
 136       // The constant value is obtained from the spv.entry_point_abi.
 137       // CHECK: spv.Constant 1 : i32
 138       %0 = "gpu.block_dim"() {dimension = "z"} : () -> index
 139       gpu.return
 140     }
 141   }
 142 }
 143
 144 // -----
 145
 146 module attributes {gpu.container_module} {
 147   func @builtin() {
 148     %c0 = arith.constant 1 : index
 149     gpu.launch_func @kernels::@builtin_local_id_x
 150         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
 151     return
 152   }
 153
 154   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
 155   // CHECK: spv.GlobalVariable [[LOCALINVOCATIONID:@.*]] built_in("LocalInvocationId")
 156   gpu.module @kernels {
 157     gpu.func @builtin_local_id_x() kernel
 158       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
 159       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[LOCALINVOCATIONID]]
 160       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
 161       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
 162       %0 = "gpu.thread_id"() {dimension = "x"} : () -> index
 163       gpu.return
 164     }
 165   }
 166 }
 167
 168 // -----
 169
 170 module attributes {gpu.container_module} {
 171   func @builtin() {
 172     %c0 = arith.constant 1 : index
 173     gpu.launch_func @kernels::@builtin_num_workgroups_x
 174         blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
 175     return
 176   }
 177
 178   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
 179   // CHECK: spv.GlobalVariable [[NUMWORKGROUPS:@.*]] built_in("NumWorkgroups")
 180   gpu.module @kernels {
 181     gpu.func @builtin_num_workgroups_x() kernel
 182       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
 183       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[NUMWORKGROUPS]]
 184       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
 185       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
 186       %0 = "gpu.grid_dim"() {dimension = "x"} : () -> index
 187       gpu.return
 188     }
 189   }
 190 }
 191
 192 // -----
 193
 194 module attributes {gpu.container_module} {
 195   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
 196   // CHECK: spv.GlobalVariable [[SUBGROUPID:@.*]] built_in("SubgroupId")
 197   gpu.module @kernels {
 198     gpu.func @builtin_subgroup_id() kernel
 199       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
 200       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[SUBGROUPID]]
 201       // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
 202       %0 = gpu.subgroup_id : index
 203       gpu.return
 204     }
 205   }
 206 }
 207
 208 // -----
 209
 210 module attributes {gpu.container_module} {
 211   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
 212   // CHECK: spv.GlobalVariable [[NUMSUBGROUPS:@.*]] built_in("NumSubgroups")
 213   gpu.module @kernels {
 214     gpu.func @builtin_num_subgroups() kernel
 215       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
 216       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[NUMSUBGROUPS]]
 217       // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
 218       %0 = gpu.num_subgroups : index
 219       gpu.return
 220     }
 221   }
 222 }
 223
 224 // -----
 225
 226 module attributes {gpu.container_module} {
 227   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
 228   // CHECK: spv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize")
 229   gpu.module @kernels {
 230     gpu.func @builtin_subgroup_size() kernel
 231       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
 232       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[SUBGROUPSIZE]]
 233       // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
 234       %0 = gpu.subgroup_size : index
 235       gpu.return
 236     }
 237   }
 238 }