save work

Author: charithaintc

mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeTranspose.cpp

@@ -12,8 +12,10 @@
 #include "mlir/Dialect/XeGPU/IR/XeGPU.h"
 #include "mlir/Dialect/XeGPU/Transforms/Passes.h"
 #include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
+#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include <optional>
 
 namespace mlir {
 namespace xegpu {
@@ -29,20 +31,36 @@ using namespace mlir;
 
 namespace {
 
-class XeGPULoadNdPattern final : public OpConversionPattern<xegpu::LoadNdOp> {
+static std::optional<SmallVector<int64_t>>
+get2DLaneData(xegpu::TensorDescType tdescType) {
+  auto layout = tdescType.getLayoutAttr();
+  if (!layout)
+    return std::nullopt;
+  auto laneData = layout.getEffectiveLaneDataAsInt();
+  if (laneData.size() != 2)
+    return std::nullopt;
+  return laneData;
+}
+
+class XeGPUCreateNdDescOpPattern final
+    : public OpConversionPattern<xegpu::CreateNdDescOp> {
 public:
-  using OpConversionPattern<xegpu::LoadNdOp>::OpConversionPattern;
+  using OpConversionPattern<xegpu::CreateNdDescOp>::OpConversionPattern;
   LogicalResult
-  matchAndRewrite(xegpu::LoadNdOp loadOp, OpAdaptor adaptor,
+  matchAndRewrite(xegpu::CreateNdDescOp createNdOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    auto tdescTy = createNdOp.getType();
+    auto convertType = this->getTypeConverter()->convertType(tdescTy);
+    if (convertType == tdescTy)
+      return failure();
     return success();
   }
 };
 } // namespace
 
 void xegpu::populateXeGPUOptimizeTransposePatterns(
     RewritePatternSet &patterns) {
-  patterns.add<XeGPULoadNdPattern>(patterns.getContext());
+  patterns.add<XeGPUCreateNdDescOpPattern>(patterns.getContext());
 }
 
 namespace {
@@ -55,6 +73,42 @@ struct XeGPUOptimizeTransposePass final
     TypeConverter converter;
     RewritePatternSet patterns(&context);
     ConversionTarget target(context);
+
+    target.addDynamicallyLegalOp<xegpu::CreateNdDescOp>(
+        [](xegpu::CreateNdDescOp createNdOp) {
+          auto optionalLaneData = get2DLaneData(createNdOp.getType());
+          if (!optionalLaneData)
+            return true;
+          auto laneData = optionalLaneData.value();
+          return laneData[0] != 1 || laneData[1] == 1;
+        });
+
+    converter.addConversion([](xegpu::TensorDescType tdescType) {
+      auto optionalLaneData = get2DLaneData(tdescType);
+      if (!optionalLaneData)
+        return tdescType;
+      auto laneData = optionalLaneData.value();
+      int64_t innerLaneData = laneData[1];
+      if (laneData[0] == 1 && innerLaneData != 1) {
+        int elementTyBitwidth =
+            tdescType.getElementType().getIntOrFloatBitWidth();
+        assert(elementTyBitwidth < 32 &&
+               "Expected element type bitwidth < 32 with laneData[1] != 1");
+        SmallVector<int64_t> newShape(tdescType.getShape());
+        newShape.back() = newShape.back() / innerLaneData;
+        Type newElemTy = IntegerType::get(tdescType.getContext(),
+                                          elementTyBitwidth * innerLaneData);
+        xegpu::LayoutAttr newLayout = xegpu::LayoutAttr::get(
+            tdescType.getContext(),
+            tdescType.getLayoutAttr().getLaneLayout().asArrayRef(), {1, 1});
+        return xegpu::TensorDescType::get(
+            newShape, newElemTy, tdescType.getArrayLength(),
+            tdescType.getBoundaryCheck(), tdescType.getMemorySpace(),
+            newLayout);
+      }
+      return tdescType;
+    });
+
     scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
                                                          target);
     xegpu::populateXeGPUOptimizeTransposePatterns(patterns);

mlir/test/Dialect/XeGPU/optimize-transpose.mlir

@@ -0,0 +1,18 @@
+func.func @gemm_b_transpose(%arg0: memref<256x256xf16>, %arg1: memref<256x256xf16>, %arg2: memref<256x256xf32>) {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %c256 = arith.constant 256 : index
+  %0 = xegpu.create_nd_tdesc %arg2 : memref<256x256xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+  %1 = xegpu.load_nd %0[%c0, %c0]  {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf32>
+  %2 = xegpu.create_nd_tdesc %arg0 : memref<256x256xf16> -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+  %3 = xegpu.create_nd_tdesc %arg1 : memref<256x256xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>>
+  %4 = scf.for %arg3 = %c0 to %c256 step %c16 iter_args(%arg4 = %1) -> (vector<8x16xf32>) {
+    %5 = xegpu.load_nd %2[%c0, %arg3]  {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
+    %6 = xegpu.load_nd %3[%c0, %arg3]  {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>> -> vector<16x16xf16>
+    %7 = vector.transpose %6, [1, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} : vector<16x16xf16> to vector<16x16xf16>
+    %8 = xegpu.dpas %5, %7, %arg4 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+    scf.yield %8 : vector<8x16xf32>
+  } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+  xegpu.store_nd %4, %0[%c0, %c0]  : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+  return
+}