generalize mechanism to find producer layout

Author: tkarna

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

@@ -209,10 +209,11 @@ def ConvertLayoutOp : Op<Transform_Dialect, "xegpu.convert_layout", [
   let summary = "Convert xegpu.layout attribute for a value.";
   let description = [{
     Adds an `xegpu.convert_layout` op to convert the `xegpu.layout` attribute
-    of a value. First, the `xegpu.load_nd` producer op of the value is found.
-    It must already be annotated with a layout. An `xegpu.convert_layout` op,
-    whose destination layout is defined by the `sg_layout`, `sg_data` and
-    optional `inst_data` attributes, is inserted after the load op.
+    of a value. The source layout is inferred by inspecting the producer ops. A
+    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.
   }];
 
   let arguments = (ins TransformValueHandleTypeInterface:$target,

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

@@ -119,6 +119,33 @@ static std::optional<T> findProducerOfType(Value val) {
   return findProducerOfType<T>(producerOp->getOperand(0));
 }
 
+/// Find layout attribute in producer chain.
+/// Traces producer ops until a layout attribute is found. Only traces through
+/// ops with a single operand, in other cases the op's result layout attribute
+/// must be set. Returns std::nullopt if no layout attribute is found.
+xegpu::LayoutAttr findProducerLayout(Value val) {
+  // Get layout attr from value or producer's attribute or operand.
+  if (auto layoutAttr = dyn_cast_if_present<xegpu::LayoutAttr>(
+          xegpu::getDistributeLayoutAttr(val)))
+    return layoutAttr;
+
+  // Recurse up the producer chain.
+  Operation *producerOp = val.getDefiningOp();
+  if (!producerOp) {
+    LDBG() << "Failed to find producer op.";
+    return nullptr;
+  }
+  if (producerOp->getNumOperands() == 0) {
+    LDBG() << "Producer has no operands.";
+    return nullptr;
+  }
+  if (producerOp->getNumOperands() > 1) {
+    LDBG() << "Producer has multiple operands.";
+    return nullptr;
+  }
+  return findProducerLayout(producerOp->getOperand(0));
+}
+
 /// Create a layout attribute from the given parameters.
 static xegpu::LayoutAttr
 createLayoutAttr(MLIRContext *ctx, ArrayRef<int32_t> sgLayout,
@@ -436,34 +463,33 @@ transform::ConvertLayoutOp::apply(transform::TransformRewriter &rewriter,
            << llvm::range_size(targetValues) << ")";
   auto value = *targetValues.begin();
 
-  xegpu::LayoutAttr layoutAttr = nullptr;
+  xegpu::LayoutAttr targetLayoutAttr = nullptr;
   auto status = getLayoutAttrFromOperands(getContext(), state, (*this),
                                           getMixedSgLayout(), getMixedSgData(),
-                                          getMixedInstData(), layoutAttr);
+                                          getMixedInstData(), targetLayoutAttr);
   if (!status.succeeded())
     return status;
 
-  // Get load op.
-  auto maybeLoadOp = findProducerOfType<xegpu::LoadNdOp>(value);
-  if (!maybeLoadOp)
-    return emitSilenceableFailure(getLoc()) << "Could not find load op.";
-  auto loadOp = *maybeLoadOp;
-  // Get load op operand value layout
-  auto producerLayoutAttr =
-      xegpu::getDistributeLayoutAttr(loadOp.getOperand(0));
+  // Find source layout attribute from the producer chain.
+  auto producerLayoutAttr = findProducerLayout(value);
   if (!producerLayoutAttr)
     return emitSilenceableFailure(getLoc())
-           << "Operand producer op does not have a layout attr.";
+           << "Could not find a layout attribute in the producer chain.";
+
+  // Find first user op to define insertion point for layout conversion.
+  if (value.use_empty())
+    return emitSilenceableFailure(getLoc())
+           << "Value has no users to insert layout conversion.";
+  Operation *userOp = *value.getUsers().begin();
 
-  if (producerLayoutAttr != layoutAttr) {
-    rewriter.setInsertionPointAfter(loadOp.getOperation());
-    auto source = loadOp.getResult();
+  if (producerLayoutAttr != targetLayoutAttr) {
+    rewriter.setInsertionPoint(userOp);
     auto convLayoutOp = xegpu::ConvertLayoutOp::create(
-        rewriter, loadOp.getLoc(), source.getType(), source, producerLayoutAttr,
-        layoutAttr);
+        rewriter, value.getLoc(), value.getType(), value, producerLayoutAttr,
+        targetLayoutAttr);
     // Replace load op result with the converted layout.
     rewriter.replaceUsesWithIf(
-        source, convLayoutOp.getResult(), [&](OpOperand &use) {
+        value, convLayoutOp.getResult(), [&](OpOperand &use) {
           return use.getOwner() != convLayoutOp.getOperation();
         });
   }

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

@@ -124,3 +124,24 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// CHECK-LABEL: @convert_layout_no_producer_attr
+func.func @convert_layout_no_producer_attr(%arg0: vector<32x32xf16>, %arg1: vector<32x32xf16>) {
+  %c0 = arith.constant 0 : index
+  %0 = arith.addf %arg0, %arg1 : vector<32x32xf16>
+  %1 = arith.extf %0 : vector<32x32xf16> to vector<32x32xf32>
+  %2 = arith.truncf %1 : vector<32x32xf32> to vector<32x32xf16>
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %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.yield
+  }
+}

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

@@ -365,9 +365,37 @@ module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
     %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 %{{.*}}
     %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.yield
   }
 }
+
+// -----
+
+// CHECK-LABEL: @convert_layout_producer_attr
+func.func @convert_layout_producer_attr(%arg0: vector<32x32xf16>, %arg1: vector<32x32xf16>) {
+  %c0 = arith.constant 0 : index
+  %0 = arith.addf %arg0, %arg1 {layout_result_0 =
+        #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>} :
+        vector<32x32xf16>
+  // CHECK: %[[V0:.+]] = arith.extf
+  %1 = arith.extf %0 : vector<32x32xf16> to vector<32x32xf32>
+  // CHECK: %[[V1:.+]] = xegpu.convert_layout %[[V0]]
+  // CHECK: input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>
+  // CHECK: target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [8, 16]>
+  // CHECK: %[[V0:.+]] = arith.truncf %[[V1]]
+  %2 = arith.truncf %1 : vector<32x32xf32> to vector<32x32xf16>
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %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.yield
+  }
+}