xegpu: remove load_data argument from set_dpas_layout transform op

Author: tkarna

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

@@ -51,23 +51,20 @@ def SetDPASLayoutOp : Op<Transform_Dialect, "xegpu.set_dpas_layout", [
     Given a `xegpu.dpas` operation, this transform adds `xegpu.layout`
     attribute to it's operand's tensor descriptor. The target operand is
     defined by the `tileIndex` argument. The layout is defined by the
-    `sg_layout`, `sg_data` and `inst_data` attributes. The `load_data`
-    attribute defines the tile size used for loading the data. It must be a
-    multiple of the `inst_data` size.
+    `sg_layout`, `sg_data` and `inst_data` attributes.
   }];
 
   let arguments = (ins TransformHandleTypeInterface : $dpasOp,
                    I64Attr : $tileIndex,
                    DenseI32ArrayAttr : $sgLayout,
                    DenseI32ArrayAttr : $sgData,
-                   OptionalAttr<DenseI32ArrayAttr> : $loadData,
                    DenseI32ArrayAttr : $instData);
 
   let results = (outs);
 
   let assemblyFormat =
       "$dpasOp `index` `=` $tileIndex `sg_layout` `=` $sgLayout `sg_data` `=` "
-      "$sgData (`load_data` `=` $loadData^)? `inst_data` `=` $instData attr-dict `:` type($dpasOp)";
+      "$sgData `inst_data` `=` $instData attr-dict `:` type($dpasOp)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(

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

@@ -610,59 +610,19 @@ DiagnosedSilenceableFailure transform::SetDPASLayoutOp::applyToOne(
            << "Expected inst_data to be a 2D vector";
   }
 
-  llvm::ArrayRef<int> loadData = instData;
-  if (getLoadData().has_value()) {
-    loadData = getLoadData().value();
-    if (loadData.size() != 2) {
-      return emitSilenceableFailure(getLoc())
-            << "Expected load_data to be a 2D vector";
-    }
-    if (loadData[0] < instData[0] || loadData[1] < instData[1]) {
-      return emitSilenceableFailure(getLoc())
-      << "load_data size must be larger or equal to inst_data size";
-    }
-    if (loadData[0] % instData[0] != 0 || loadData[1] % instData[1] != 0) {
-      return emitSilenceableFailure(getLoc())
-      << "load_data must be evenly divisible by inst_data";
-    }
-  }
-
   // Replace descriptor op using layout attribute.
   Value opVec = dpasOp.getOperation()->getOperand(tileIndex);
   auto maybeDescOp = findDescriptorOp(opVec, dpasOp.getOperation());
   if (!maybeDescOp) {
     return emitSilenceableFailure(getLoc()) << "Could not find descriptor op.";
   }
   auto descOp = *maybeDescOp;
-  // Layout for the load op.
-  auto loadLayoutAttr = createLayoutAttr(rewriter.getContext(), sgLayout, sgData, loadData);
-  descOp = setDescLayout(rewriter, descOp, loadLayoutAttr);
-  // Layout for the instruction.
-  auto instLayoutAttr = createLayoutAttr(rewriter.getContext(), sgLayout, sgData, instData);
+  // Set layout attribute.
+  auto layoutAttr = createLayoutAttr(rewriter.getContext(), sgLayout, sgData, instData);
+  descOp = setDescLayout(rewriter, descOp, layoutAttr);
   if (tileIndex == 2) {
-    // C operand: set layout attribute for the dpas op result
-    xegpu::setLayoutAttr(dpasOp.getOperation()->getResults()[0], instLayoutAttr);
-  }
-
-  if (loadLayoutAttr != instLayoutAttr) {
-    // Insert convert layout op after load op.
-    auto maybeLoadOp = getUserOfType<xegpu::LoadNdOp>(descOp.getResult());
-    if (!maybeLoadOp) {
-      return emitSilenceableFailure(getLoc())
-      << "Expected a xegpu.load_nd op as a user of the descriptor op."; 
-    }
-    auto loadOp = *maybeLoadOp;
-    rewriter.setInsertionPointAfter(loadOp.getOperation());
-    auto source = loadOp.getResult();
-    auto convLayoutOp = rewriter.create<xegpu::ConvertLayoutOp>(
-        loadOp.getLoc(), source.getType(), source,
-        loadLayoutAttr, instLayoutAttr);
-    // Replace load op result with the converted layout.
-    rewriter.replaceUsesWithIf(
-      source, convLayoutOp.getResult(),
-      [&](OpOperand &use) {
-        return use.getOwner() != convLayoutOp.getOperation();
-      });
+    // C operand: set layout attribute for the dpas op result.
+    xegpu::setLayoutAttr(dpasOp.getOperation()->getResults()[0], layoutAttr);
   }
 
   return DiagnosedSilenceableFailure::success();

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

@@ -28,7 +28,6 @@ def __init__(
         sg_data: Union[Sequence[int], Attribute],
         inst_data: Union[Sequence[int], Attribute],
         *,
-        load_data: Optional[Union[Sequence[int], Attribute]] = None,
         loc=None,
         ip=None,
     ):
@@ -38,7 +37,6 @@ def __init__(
             sg_layout,
             sg_data,
             inst_data,
-            loadData=load_data,
             loc=loc,
             ip=ip
         )

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

@@ -107,37 +107,7 @@ 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
     // CHECK: transform.xegpu.set_dpas_layout %{{.*}}
-    transform.xegpu.set_dpas_layout %0 index = 2 sg_layout = [8, 4] sg_data = [32, 64] load_data = [8, 16] inst_data = [8, 16] : !transform.any_op
-    transform.yield
-  }
-}
-
-// -----
-
-// CHECK-LABEL: @set_dpas_layout_load_a
-func.func @set_dpas_layout_load_a(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
-  // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0
-  // CHECK-SAME: #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>>
-  %0 = xegpu.create_nd_tdesc %arg0[0, 0] : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
-  // CHECK: %[[V1:.+]] = xegpu.load_nd %[[V0]]
-  // CHECK: %[[V2:.+]] = xegpu.convert_layout %[[V1]]
-  // CHECK-SAME: resMap = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [8, 16]>
-  // CHECK-SAME: srcMap = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>
-  %1 = xegpu.load_nd %0  : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
-  %2 = xegpu.create_nd_tdesc %arg1[0, 0] : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>
-  %3 = xegpu.load_nd %2  : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>
-  %4 = xegpu.create_nd_tdesc %arg2[0, 0] : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
-  %5 = xegpu.load_nd %4  : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
-  // CHECK: = xegpu.dpas %[[V2]]
-  %6 = xegpu.dpas %1, %3, %5 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
-  return
-}
-
-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
-    // CHECK: transform.xegpu.set_dpas_layout %{{.*}}
-    transform.xegpu.set_dpas_layout %0 index = 0 sg_layout = [8, 4] sg_data = [32, 32] load_data = [32, 16] inst_data = [8, 16] : !transform.any_op
+    transform.xegpu.set_dpas_layout %0 index = 2 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op
     transform.yield
   }
 }