convert_layout transform op returns a handle to created convert op

Author: tkarna

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

@@ -257,7 +257,8 @@ def ConvertLayoutOp : Op<Transform_Dialect, "xegpu.convert_layout", [
     failure is emitted if source layout cannot be found. An
     `xegpu.convert_layout` op, whose destination layout is defined by the
     `sg_layout`, `sg_data` and optional `inst_data` attributes, is emitted
-    before the first use of the value.
+    before the first use of the value. Returns a handle to the emitted
+    `xegpu.convert_layout` op.
   }];
 
   let arguments = (ins TransformValueHandleTypeInterface:$target,
@@ -269,7 +270,7 @@ def ConvertLayoutOp : Op<Transform_Dialect, "xegpu.convert_layout", [
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_inst_data
                    );
 
-  let results = (outs);
+  let results = (outs TransformHandleTypeInterface:$newConvertOp);
   let builders = [
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedSgLayout,
@@ -283,7 +284,7 @@ def ConvertLayoutOp : Op<Transform_Dialect, "xegpu.convert_layout", [
     `sg_layout` `=` custom<DynamicIndexList>($sg_layout, $static_sg_layout)
     `sg_data` `=` custom<DynamicIndexList>($sg_data, $static_sg_data)
     (`inst_data` `=` custom<DynamicIndexList>($inst_data, $static_inst_data)^)?
-    attr-dict `:` qualified(type(operands))
+    attr-dict `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{

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

@@ -614,18 +614,18 @@ transform::ConvertLayoutOp::apply(transform::TransformRewriter &rewriter,
            << "Value has no users to insert layout conversion.";
   Operation *userOp = *value.getUsers().begin();
 
-  if (producerLayoutAttr != targetLayoutAttr) {
-    rewriter.setInsertionPoint(userOp);
-    auto convLayoutOp = xegpu::ConvertLayoutOp::create(
-        rewriter, value.getLoc(), value.getType(), value, producerLayoutAttr,
-        targetLayoutAttr);
-    // Replace load op result with the converted layout.
-    rewriter.replaceUsesWithIf(
-        value, convLayoutOp.getResult(), [&](OpOperand &use) {
-          return use.getOwner() != convLayoutOp.getOperation();
-        });
-  }
-
+  // Emit convert_layout op.
+  rewriter.setInsertionPoint(userOp);
+  auto convLayoutOp = xegpu::ConvertLayoutOp::create(
+      rewriter, value.getLoc(), value.getType(), value, producerLayoutAttr,
+      targetLayoutAttr);
+  // Replace load op result with the converted layout.
+  rewriter.replaceUsesWithIf(
+      value, convLayoutOp.getResult(), [&](OpOperand &use) {
+        return use.getOwner() != convLayoutOp.getOperation();
+      });
+
+  results.set(llvm::cast<OpResult>(getResult()), {convLayoutOp});
   return DiagnosedSilenceableFailure::success();
 }
 
@@ -635,6 +635,7 @@ void transform::ConvertLayoutOp::getEffects(
   onlyReadsHandle(getSgLayoutMutable(), effects);
   onlyReadsHandle(getSgDataMutable(), effects);
   onlyReadsHandle(getInstDataMutable(), effects);
+  producesHandle(getOperation()->getOpResults(), effects);
   modifiesPayload(effects);
 }
 

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

@@ -295,6 +295,7 @@ def __init__(
             _,
         ) = _dispatch_dynamic_index_list(inst_data)
         super().__init__(
+            transform.AnyOpType.get(),
             target,
             dynamic_sg_layout,
             dynamic_sg_data,
@@ -323,4 +324,4 @@ def convert_layout(
         inst_data=inst_data,
         loc=loc,
         ip=ip,
-    )
+    ).result

mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir

@@ -172,7 +172,7 @@ module attributes {transform.with_named_sequence} {
     %0 = transform.structured.match ops{["arith.truncf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
     %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
     // expected-error@below {{Could not find a layout attribute in the producer chain.}}
-    transform.xegpu.convert_layout %1 sg_layout = [8, 4] sg_data = [32, 32] inst_data = [8, 16] : !transform.any_value
+    transform.xegpu.convert_layout %1 sg_layout = [8, 4] sg_data = [32, 32] inst_data = [8, 16] : (!transform.any_value) -> !transform.any_op
     transform.yield
   }
 }

mlir/test/Dialect/XeGPU/transform-ops.mlir

@@ -427,7 +427,7 @@ module attributes {transform.with_named_sequence} {
     %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg1 : (!transform.any_op) -> !transform.any_op
     %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
     // CHECK: transform.xegpu.convert_layout %{{.*}}
-    transform.xegpu.convert_layout %1 sg_layout = [8, 4] sg_data = [32, 32] inst_data = [8, 16] : !transform.any_value
+    transform.xegpu.convert_layout %1 sg_layout = [8, 4] sg_data = [32, 32] inst_data = [8, 16] : (!transform.any_value) -> !transform.any_op
     transform.yield
   }
 }
@@ -459,7 +459,7 @@ module attributes {transform.with_named_sequence} {
     %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
     %layout0 = transform.param.constant 8 : i64 -> !transform.param<i64>
     // CHECK: transform.xegpu.convert_layout %{{.*}}
-    transform.xegpu.convert_layout %1 sg_layout = [%layout0, 4] sg_data = [32, 32] inst_data = [8, 16] : !transform.any_value, !transform.param<i64>
+    transform.xegpu.convert_layout %1 sg_layout = [%layout0, 4] sg_data = [32, 32] inst_data = [8, 16] : (!transform.any_value, !transform.param<i64>) -> !transform.any_op
     transform.yield
   }
 }
@@ -487,7 +487,7 @@ module attributes {transform.with_named_sequence} {
     %0 = transform.structured.match ops{["arith.truncf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
     %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
     // CHECK: transform.xegpu.convert_layout %{{.*}}
-    transform.xegpu.convert_layout %1 sg_layout = [8, 4] sg_data = [32, 32] inst_data = [8, 16] : !transform.any_value
+    transform.xegpu.convert_layout %1 sg_layout = [8, 4] sg_data = [32, 32] inst_data = [8, 16] : (!transform.any_value) -> !transform.any_op
     transform.yield
   }
 }