[Mlir-commits] [mlir] [MLIR][XeGPU][TransformOps] Add set_gpu_launch_threads op (PR #166865)

[Rolf Morel]

================
@@ -132,3 +132,39 @@ def __init__(
             loc=loc,
             ip=ip,
         )
+
+
+class SetGPULaunchThreadsOp(SetGPULaunchThreadsOp):
+    """Specialization for SetGPULaunchThreadsOp class."""
+
+    def __init__(
+        self,
+        launch_op: Union[Operation, Value],
+        threads: MixedValues,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        (
+            dynamic_threads,
+            static_threads,
+            _,
+        ) = _dispatch_dynamic_index_list(threads)
+
+        super().__init__(
+            _get_op_result_or_value(launch_op),
+            dynamic_threads,
+            static_threads=static_threads,
+            loc=loc,
+            ip=ip,
+        )
+
+
+def set_gpu_launch_threads(
+    launch_op: Union[Operation, Value],
+    threads: MixedValues,
+    *,
+    loc=None,
+    ip=None,
+) -> SetGPULaunchThreadsOp:
+    return SetGPULaunchThreadsOp(launch_op, threads, loc=loc, ip=ip)
----------------
rolfmorel wrote:

Yes!

(I forgot to mention: in case the op returns results, the snake_case wrapper should return the result/results and only return the op directly if the op doesn't have results. This is what the autogenerated bindings do now. For examples, see: https://github.com/llvm/llvm-project/pull/166871/files#diff-fecd4bdb8ee395146efaf323769aef391fcb3474f054100e63b01f4e4f52f6d0) 

https://github.com/llvm/llvm-project/pull/166865