[mlir][gpu]add AffineScope Trait to gpu.launch. (#121058)

add AffineScope Trait to gpu.launch.

Author: linuxlonelyeagle

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<InferIntRangeInterface, ["inferResultRanges"]>,
-      RecursiveMemoryEffects]>,
+      GPU_AsyncOpInterface, RecursiveMemoryEffects]>,
     Arguments<(ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
                Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,

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<?x?xf32>) kernel {
-      %c3 = arith.constant 1 : index
-      %dim = memref.dim %arg0, %c3 : memref<?x?xf32>
-      %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<?x?xf32>) kernel {
-// CHECK:             %[[VAL_1:.*]] = arith.constant 1 : index
-// CHECK:             %[[VAL_2:.*]] = memref.dim %[[VAL_0]], %[[VAL_1]] : memref<?x?xf32>
-// 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 {
 
 // -----
 

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(