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,
@@ -92,11 +92,11 @@ def SetOpLayoutAttrOp : Op<Transform_Dialect, "xegpu.set_op_layout_attr", [
     is defined by the `sg_layout`, `sg_data` and optional `inst_data` attributes.
   }];
 
-  let arguments = (ins TransformHandleTypeInterface : $target,
-                   DefaultValuedOptionalAttr<I64Attr, "0"> : $index,
-                   Variadic<TransformAnyParamTypeOrAnyHandle> : $sg_layout,
-                   Variadic<TransformAnyParamTypeOrAnyHandle> : $sg_data,
-                   Variadic<TransformAnyParamTypeOrAnyHandle> : $inst_data,
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                   DefaultValuedOptionalAttr<I64Attr, "0">:$index,
+                   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,

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

@@ -92,8 +92,7 @@ createLayoutAttr(MLIRContext *ctx, ArrayRef<int32_t> sgLayout,
 
 /// Generate `xegpu::LayoutAttr` from op mixed layout values.
 DiagnosedSilenceableFailure
-getLayoutAttrFromOperands(transform::TransformRewriter &rewriter,
-                          transform::TransformState &state,
+getLayoutAttrFromOperands(MLIRContext *ctx, transform::TransformState &state,
                           TransformOpInterface transformOp,
                           ArrayRef<::mlir::OpFoldResult> mixedSgLayout,
                           ArrayRef<::mlir::OpFoldResult> mixedSgData,
@@ -116,8 +115,7 @@ getLayoutAttrFromOperands(transform::TransformRewriter &rewriter,
                            ? std::nullopt
                            : std::optional<ArrayRef<int32_t>>(instData);
 
-  layoutAttr =
-      createLayoutAttr(rewriter.getContext(), sgLayout, sgData, maybeInstData);
+  layoutAttr = createLayoutAttr(ctx, sgLayout, sgData, maybeInstData);
 
   return DiagnosedSilenceableFailure::success();
 }
@@ -175,7 +173,7 @@ transform::SetDescLayoutOp::apply(transform::TransformRewriter &rewriter,
   Operation *target = *targetOps.begin();
 
   xegpu::LayoutAttr layoutAttr = nullptr;
-  auto status = getLayoutAttrFromOperands(rewriter, state, (*this),
+  auto status = getLayoutAttrFromOperands(getContext(), state, (*this),
                                           getMixedSgLayout(), getMixedSgData(),
                                           getMixedInstData(), layoutAttr);
   if (!status.succeeded())
@@ -235,7 +233,6 @@ DiagnosedSilenceableFailure
 transform::SetOpLayoutAttrOp::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 "
@@ -256,18 +253,17 @@ transform::SetOpLayoutAttrOp::apply(transform::TransformRewriter &rewriter,
   }
 
   xegpu::LayoutAttr layoutAttr = nullptr;
-  auto status = getLayoutAttrFromOperands(rewriter, state, (*this),
+  auto status = getLayoutAttrFromOperands(getContext(), state, (*this),
                                           getMixedSgLayout(), getMixedSgData(),
                                           getMixedInstData(), layoutAttr);
   if (!status.succeeded())
     return status;
 
   // Set layout attribute for the op result or operand
-  if (resultTarget) {
+  if (resultTarget)
     xegpu::setDistributeLayoutAttr(target->getResult(index), layoutAttr);
-  } else {
+  else
     xegpu::setDistributeLayoutAttr(target->getOpOperand(index), layoutAttr);
-  }
   return DiagnosedSilenceableFailure::success();
 }
 

mlir/python/mlir/dialects/transform/xegpu.py

@@ -76,9 +76,9 @@ def __init__(
         sg_layout: MixedValues,
         sg_data: MixedValues,
         *,
-        inst_data: MixedValues = None,
-        index: Union[int, Attribute] = None,
-        result: Union[bool, Attribute] = None,
+        inst_data: Optional[MixedValues] = None,
+        index: Optional[Union[int, Attribute]] = None,
+        result: Optional[Union[bool, Attribute]] = None,
         loc=None,
         ip=None,
     ):