address review comments

Author: tkarna

mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td

@@ -31,16 +31,16 @@ def SetDescLayoutOp : Op<Transform_Dialect, "xegpu.set_desc_layout", [
   }];
 
   let arguments = (ins
-                   TransformHandleTypeInterface : $target,
-                   Variadic<TransformAnyParamTypeOrAnyHandle> : $sg_layout,
-                   Variadic<TransformAnyParamTypeOrAnyHandle> : $sg_data,
-                   Variadic<TransformAnyParamTypeOrAnyHandle> : $inst_data,
+                   TransformHandleTypeInterface:$target,
+                   Variadic<TransformAnyParamTypeOrAnyHandle>:$sg_layout,
+                   Variadic<TransformAnyParamTypeOrAnyHandle>:$sg_data,
+                   Variadic<TransformAnyParamTypeOrAnyHandle>:$inst_data,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sg_layout,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sg_data,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_inst_data
                    );
 
-  let results = (outs TransformHandleTypeInterface : $transformed);
+  let results = (outs TransformHandleTypeInterface:$transformed);
   let builders = [
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedSgLayout,
@@ -84,11 +84,13 @@ def SetGPULaunchThreadsOp
       TransformOpInterface
     ]> {
 
-  let summary = "Set number of threads for a given gpu.launch operation";
-  let description = "Set number of threads for a given `gpu.launch` operation.";
+  let summary = "Set number of threads for a given gpu.launch operation"; let
+  description = [{
+    Overrides the x,y,z threads operands of a given `gpu.launch` operation in-place.
+  }];
 
-  let arguments = (ins TransformHandleTypeInterface : $target,
-                   Variadic<TransformAnyParamTypeOrAnyHandle> : $threads,
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                   Variadic<TransformAnyParamTypeOrAnyHandle>:$threads,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_threads
                    );
   let results = (outs);

mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp

@@ -210,7 +210,6 @@ DiagnosedSilenceableFailure
 transform::SetGPULaunchThreadsOp::apply(transform::TransformRewriter &rewriter,
                                         transform::TransformResults &results,
                                         transform::TransformState &state) {
-
   auto targetOps = state.getPayloadOps(getTarget());
   if (!llvm::hasSingleElement(targetOps)) {
     return emitDefiniteFailure() << "Requires exactly one targetOp handle (got "
@@ -234,7 +233,8 @@ transform::SetGPULaunchThreadsOp::apply(transform::TransformRewriter &rewriter,
 
   if (threads.size() != 3) {
     return emitSilenceableFailure(getLoc())
-           << "Expected threads to be a 3D vector";
+           << "Expected threads argument to consist of three values (got "
+           << threads.size() << ")";
   }
 
   rewriter.setInsertionPoint(launchOp);

mlir/test/python/dialects/transform_xegpu_ext.py

@@ -56,7 +56,7 @@ def setGPULaunchThreadsOp():
     sequence = transform.SequenceOp(
         transform.FailurePropagationMode.Propagate,
         [],
-        transform.OperationType.get("gpu.lauch"),
+        transform.OperationType.get("gpu.launch"),
     )
     with InsertionPoint(sequence.body):
         xegpu.SetGPULaunchThreadsOp(sequence.bodyTarget, threads=[8, 4, 1])