diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td index 3adfd5f4f2c43..2b1ce573effd0 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -794,9 +794,9 @@ def GPU_LaunchFuncOp :GPU_Op<"launch_func", [ } def GPU_LaunchOp : GPU_Op<"launch", [ - AutomaticAllocationScope, AttrSizedOperandSegments, GPU_AsyncOpInterface, + AffineScope, AutomaticAllocationScope, AttrSizedOperandSegments, DeclareOpInterfaceMethods, - RecursiveMemoryEffects]>, + GPU_AsyncOpInterface, RecursiveMemoryEffects]>, Arguments<(ins Variadic:$asyncDependencies, Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ, Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ, diff --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir index e3721806989bb..233c18c82831c 100644 --- a/mlir/test/Dialect/Affine/ops.mlir +++ b/mlir/test/Dialect/Affine/ops.mlir @@ -301,29 +301,29 @@ func.func @linearize_mixed(%index0: index, %index1: index, %index2: index, %basi // ----- -#map = affine_map<()[s0] -> (s0)> +// CHECK-LABEL: @gpu_launch_affine -// CHECK-LABEL: @gpu_affine_for +// Test `thread_id` in AffineScope, the `thread_id` is in AffineScope's toplevel, +// it is a valid symbol. -module attributes {gpu.container_module} { - gpu.module @gpu { - gpu.func @gpu_affine_for(%arg0: memref) kernel { - %c3 = arith.constant 1 : index - %dim = memref.dim %arg0, %c3 : memref - %c0 = arith.constant 0 : index - affine.for %arg3 = %c0 to #map()[%dim] step 32 { +module { + func.func @gpu_launch_affine() { + %c1 = arith.constant 1 : index + gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1) + threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) { + %thread_id_x = gpu.thread_id x + %c128 = arith.constant 128 : index + affine.for %arg12 = %thread_id_x to %c128 step 8 { } - gpu.return + gpu.terminator } + return } } -// CHECK-SAME: (%[[VAL_0:.*]]: memref) kernel { -// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index -// CHECK: %[[VAL_2:.*]] = memref.dim %[[VAL_0]], %[[VAL_1]] : memref -// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index -// CHECK: affine.for %[[VAL_4:.*]] = %[[VAL_3]] to %[[VAL_2]] step 32 { -// CHECK: } -// CHECK: gpu.return + +// CHECK: %[[THREAD_ID:.*]] = gpu.thread_id x +// CHECK: %[[VAL:.*]] = arith.constant 128 : index +// CHECK: affine.for %{{.*}} = %[[THREAD_ID]] to %[[VAL]] step 8 { // ----- diff --git a/mlir/test/Dialect/GPU/transform-gpu.mlir b/mlir/test/Dialect/GPU/transform-gpu.mlir index 0a5c85336831a..09ae0f4af686f 100644 --- a/mlir/test/Dialect/GPU/transform-gpu.mlir +++ b/mlir/test/Dialect/GPU/transform-gpu.mlir @@ -43,7 +43,7 @@ module attributes {transform.with_named_sequence} { !type = memref<2 x 32 x f32> !type1d = memref<32 x f32> -// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 floordiv 128)> +// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 floordiv 128)> // CHECK-LABEL: func.func @warpgroup_3d( // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> @@ -647,7 +647,7 @@ module attributes {transform.with_named_sequence} { #map = affine_map<(d0) -> (d0 * 128)> #map1 = affine_map<(d0) -> (d0 * 32)> -// CHECK-DAG: #[[$MAPB:.*]] = affine_map<()[s0] -> (s0 * 128)> +// CHECK-DAG: #[[$MAPB:.*]] = affine_map<()[s0] -> (s0 * 128)> // CHECK-DAG: #[[$MAPW:.*]] = affine_map<()[s0, s1, s2] -> (s2 * 32 + ((s0 + s1 * 4) floordiv 32) * 32)> // CHECK-LABEL: func.func @simple_fill(