more nit comments

Author: tkarna

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

@@ -7,7 +7,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.h"
-#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/XeGPU/IR/XeGPU.h"
 #include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
@@ -34,7 +33,6 @@ class XeGPUTransformDialectExtension
 void XeGPUTransformDialectExtension::init() {
   declareGeneratedDialect<scf::SCFDialect>();
   declareGeneratedDialect<arith::ArithDialect>();
-  declareGeneratedDialect<gpu::GPUDialect>();
   declareGeneratedDialect<xegpu::XeGPUDialect>();
 
   registerTransformOps<
@@ -173,33 +171,30 @@ DiagnosedSilenceableFailure
 transform::SetDescLayoutOp::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 "
                                  << llvm::range_size(targetOps) << ")";
   }
   Operation *target = *targetOps.begin();
 
-  auto transformOp = cast<TransformOpInterface>(getOperation());
-
   SmallVector<int32_t> sgLayout;
   DiagnosedSilenceableFailure status =
-      convertMixedValuesToInt(state, transformOp, sgLayout, getMixedSgLayout());
+      convertMixedValuesToInt(state, (*this), sgLayout, getMixedSgLayout());
   if (!status.succeeded())
     return status;
 
   SmallVector<int32_t> sgData;
   status =
-      convertMixedValuesToInt(state, transformOp, sgData, getMixedSgData());
+      convertMixedValuesToInt(state, (*this), sgData, getMixedSgData());
   if (!status.succeeded())
     return status;
   auto maybeSgData =
       sgData.empty() ? std::nullopt : std::optional<ArrayRef<int32_t>>(sgData);
 
   SmallVector<int32_t> instData;
   status =
-      convertMixedValuesToInt(state, transformOp, instData, getMixedInstData());
+      convertMixedValuesToInt(state, (*this), instData, getMixedInstData());
   if (!status.succeeded())
     return status;
   auto maybeInstData = instData.empty()

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

@@ -17,7 +17,7 @@
 except ImportError as e:
     raise RuntimeError("Error loading imports from extension module") from e
 
-from typing import Union
+from typing import Union, Optional
 
 
 @_ods_cext.register_operation(_Dialect, replace=True)
@@ -29,8 +29,8 @@ def __init__(
         target: Union[Operation, Value],
         sg_layout: MixedValues,
         *,
-        sg_data: MixedValues = None,
-        inst_data: MixedValues = None,
+        sg_data: Optional[MixedValues] = None,
+        inst_data: Optional[MixedValues] = None,
         loc=None,
         ip=None,
     ):