refactor ConvertLayoutPattern for wg to sg.

Author: chencha3

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

@@ -313,13 +313,13 @@ LogicalResult TensorDescType::verify(
   if (rank != 1 && rank != 2)
     return emitError() << "expected 1D or 2D tensor";
 
-  // auto blockAttr = mlir::dyn_cast_if_present<BlockTensorDescAttr>(encoding);
-  // if (blockAttr) {
-  //   MemorySpaceAttr memorySpaceAttr = blockAttr.getMemorySpace();
-  //   if (rank == 2 && memorySpaceAttr &&
-  //       memorySpaceAttr.getValue() == MemorySpace::SLM)
-  //     return emitError() << "SLM is not supported for 2D block tensor";
-  // }
+  auto blockAttr = mlir::dyn_cast_if_present<BlockTensorDescAttr>(encoding);
+  if (blockAttr) {
+    MemorySpaceAttr memorySpaceAttr = blockAttr.getMemorySpace();
+    if (rank == 2 && memorySpaceAttr &&
+        memorySpaceAttr.getValue() == MemorySpace::SLM)
+      return emitError() << "SLM is not supported for 2D block tensor";
+  }
 
   // for gather and scatter ops, Low-precision types are packed in 32-bit units.
   unsigned bitWidth = elementType.getIntOrFloatBitWidth();

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

@@ -616,7 +616,7 @@ LogicalResult ConvertLayoutOp::verify() {
   if (!resLayout)
     return emitOpError("expected target layout.");
 
-  // both srcMap and resMap should be WgLayout or SgLayout at the same time.
+  // both input and target layouts should be WgLayout or SgLayout at the same time.
   if ((!srcLayout.isWgLayout() || !resLayout.isWgLayout()) &&
       (!srcLayout.isSgLayout() || !resLayout.isSgLayout()))
     return emitOpError("expected input layout and target layout be WgLayout or "
@@ -644,10 +644,11 @@ struct FoldConvertLayoutOp : public OpRewritePattern<xegpu::ConvertLayoutOp> {
   using OpRewritePattern<xegpu::ConvertLayoutOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(xegpu::ConvertLayoutOp op,
                                 PatternRewriter &rewriter) const override {
-    if (op.getInputLayout() != op.getTargetLayout())
-      return failure();
-    rewriter.replaceOp(op, op.getSource());
-    return success();
+    if (op.getInputLayout() == op.getTargetLayout()) {
+      rewriter.replaceOp(op, op.getSource());
+      return success();
+    }
+    return failure();
   }
 };
 

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

@@ -57,39 +57,6 @@ getSgShapeAndCount(ArrayRef<int64_t> shape, xegpu::LayoutAttr layout) {
   return std::make_pair(sgShape, count);
 }
 
-// Calculate offset for each subgroup
-static SmallVector<OpFoldResult>
-calculateGlobalOffsets(ConversionPatternRewriter &rewriter, Location loc,
-                       const SmallVector<OpFoldResult> &originalOffsets,
-                       const SmallVector<Value> &localOffset,
-                       const SmallVector<int64_t> &distUnitBaseAddr,
-                       const SmallVector<int64_t> &distUnitShape) {
-  assert(localOffset.size() == distUnitBaseAddr.size() &&
-         "localOffset and distUnitBaseAddr must have the same rank");
-
-  SmallVector<OpFoldResult> globalOffsets(originalOffsets.begin(),
-                                          originalOffsets.end());
-  size_t rank = localOffset.size();
-  for (size_t i = 0; i < rank; ++i) {
-    size_t dimIdx = originalOffsets.size() - rank + i;
-    Value constOffset =
-        rewriter.create<arith::ConstantIndexOp>(loc, distUnitBaseAddr[i]);
-    Value offset =
-        rewriter.createOrFold<index::AddOp>(loc, localOffset[i], constOffset);
-    Value modValue =
-        rewriter.create<arith::ConstantIndexOp>(loc, distUnitShape[i]);
-    Value offsetMod =
-        rewriter.createOrFold<index::RemUOp>(loc, offset, modValue);
-    Value origOffset =
-        getValueOrCreateConstantIndexOp(rewriter, loc, originalOffsets[dimIdx]);
-    Value globalOffset =
-        rewriter.createOrFold<index::AddOp>(loc, origOffset, offsetMod);
-    globalOffsets[dimIdx] = globalOffset;
-  }
-
-  return globalOffsets;
-}
-
 /// This pattern transforms the CreateNdDescOp to create a subgroup descriptor
 /// from a workgroup descriptor. It replaces the offsets and sizes with
 /// appropriate values for the subgroup.
@@ -138,6 +105,39 @@ calculateGlobalOffsets(ConversionPatternRewriter &rewriter, Location loc,
 struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
   using OpConversionPattern<xegpu::CreateNdDescOp>::OpConversionPattern;
 
+  // Calculate offset for each subgroup
+  static SmallVector<OpFoldResult>
+  calculateGlobalOffsets(ConversionPatternRewriter &rewriter, Location loc,
+                         const SmallVector<OpFoldResult> &originalOffsets,
+                         const SmallVector<Value> &localOffset,
+                         const SmallVector<int64_t> &distUnitBaseAddr,
+                         const SmallVector<int64_t> &distUnitShape) {
+    assert(localOffset.size() == distUnitBaseAddr.size() &&
+           "localOffset and distUnitBaseAddr must have the same rank");
+
+    SmallVector<OpFoldResult> globalOffsets(originalOffsets.begin(),
+                                            originalOffsets.end());
+    size_t rank = localOffset.size();
+    for (size_t i = 0; i < rank; ++i) {
+      size_t dimIdx = originalOffsets.size() - rank + i;
+      Value constOffset =
+          rewriter.create<arith::ConstantIndexOp>(loc, distUnitBaseAddr[i]);
+      Value offset =
+          rewriter.createOrFold<index::AddOp>(loc, localOffset[i], constOffset);
+      Value modValue =
+          rewriter.create<arith::ConstantIndexOp>(loc, distUnitShape[i]);
+      Value offsetMod =
+          rewriter.createOrFold<index::RemUOp>(loc, offset, modValue);
+      Value origOffset =
+          getValueOrCreateConstantIndexOp(rewriter, loc, originalOffsets[dimIdx]);
+      Value globalOffset =
+          rewriter.createOrFold<index::AddOp>(loc, origOffset, offsetMod);
+      globalOffsets[dimIdx] = globalOffset;
+    }
+
+    return globalOffsets;
+  }
+
   LogicalResult
   matchAndRewrite(xegpu::CreateNdDescOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
@@ -390,21 +390,6 @@ struct WgToSgElementwiseOp : public ConversionPattern {
   }
 };
 
-// based on the size of the given vector type
-static TypedValue<MemRefType>
-allocateSLMBuffer(ConversionPatternRewriter &rewriter, Location loc,
-                  VectorType type) {
-  int64_t bits = type.getElementType().getIntOrFloatBitWidth();
-  int64_t slmSizeInBytes = type.getNumElements() * bits / 8;
-  auto slmTy = MemRefType::get(slmSizeInBytes, rewriter.getI8Type(), {}, 3);
-  auto slm = rewriter.create<memref::AllocOp>(loc, slmTy);
-  auto viewTy = MemRefType::get(type.getShape(), type.getElementType(), {}, 3);
-  auto view = rewriter.create<memref::ViewOp>(
-      loc, viewTy, slm, rewriter.create<arith::ConstantIndexOp>(loc, 0),
-      ValueRange());
-  return view;
-}
-
 struct WgToSgConvertLayoutOp
     : public OpConversionPattern<xegpu::ConvertLayoutOp> {
   using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
@@ -418,115 +403,29 @@ struct WgToSgConvertLayoutOp
       return rewriter.notifyMatchFailure(
           op, "Input and target layouts must have subgroup layout");
 
-    // initialize values with the source values
-    SmallVector<Value> values(adaptor.getSource());
-
-    Location loc = op.getLoc();
-    MLIRContext *ctx = op.getContext();
-    VectorType type = op.getResult().getType();
-    ArrayRef<int64_t> shape = type.getShape();
-
     DenseI32ArrayAttr inputSgLayout = input.getSgLayout();
     DenseI32ArrayAttr inputSgData = input.getSgData();
     DenseI32ArrayAttr targetSgLayout = target.getSgLayout();
     DenseI32ArrayAttr targetSgData = target.getSgData();
 
-    // we only need SLM support when input and target layouts are different
-    if (inputSgLayout != targetSgLayout || inputSgData != targetSgData) {
-      values.clear();
-      rewriter.setInsertionPoint(op);
-      TypedValue<MemRefType> slmBuffer = allocateSLMBuffer(rewriter, loc, type);
-
-      auto linearSgId = rewriter.create<gpu::SubgroupIdOp>(
-          loc, rewriter.getIndexType(), nullptr);
-
-      { // store to slm buffer
-        SmallVector<int64_t> sgLayout =
-            llvm::to_vector_of<int64_t>(input.getSgLayout().asArrayRef());
-        SmallVector<int64_t> sgShape = getSgShapeAndCount(shape, input).first;
-        auto delinearized = affine::delinearizeIndex(
-            rewriter, loc, linearSgId, getAsIndexOpFoldResult(ctx, sgLayout));
-        if (failed(delinearized))
-          return rewriter.notifyMatchFailure(op, "Failed to delinearize sgId");
-        SmallVector<Value> sgIds = *delinearized;
-
-        SmallVector<int64_t> distUnitShape(sgLayout.size());
-        SmallVector<Value> localOffset(sgLayout.size());
-        for (size_t i = 0; i < sgLayout.size(); i++) {
-          distUnitShape[i] = std::min(sgLayout[i] * sgShape[i], shape[i]);
-          localOffset[i] = rewriter.createOrFold<index::MulOp>(
-              loc, sgIds[i],
-              rewriter.create<arith::ConstantIndexOp>(loc, sgShape[i]));
-        }
-
-        auto tdescTy = xegpu::TensorDescType::get(
-            sgShape, type.getElementType(), 1, false, xegpu::MemorySpace::SLM,
-            input.dropSgLayoutAndData());
-
-        SmallVector<OpFoldResult> zeros = getAsIndexOpFoldResult(
-            ctx, SmallVector<int64_t>(sgLayout.size(), 0));
-        for (auto [data, baseOffsets] :
-             llvm::zip_equal(adaptor.getSource(),
-                             StaticTileOffsetRange(shape, distUnitShape))) {
-          SmallVector<OpFoldResult> offsets = calculateGlobalOffsets(
-              rewriter, loc, zeros, localOffset, baseOffsets, distUnitShape);
-          auto tdesc = rewriter.create<xegpu::CreateNdDescOp>(
-              loc, tdescTy, slmBuffer, offsets);
-          rewriter.create<xegpu::StoreNdOp>(loc, data, tdesc, nullptr, nullptr,
-                                            nullptr);
-        }
-      }
-
-      rewriter.create<gpu::BarrierOp>(loc);
-
-      { // load from SLM
-        SmallVector<int64_t> sgLayout =
-            llvm::to_vector_of<int64_t>(target.getSgLayout().asArrayRef());
-        SmallVector<int64_t> sgShape = getSgShapeAndCount(shape, target).first;
-        auto delinearized = affine::delinearizeIndex(
-            rewriter, loc, linearSgId, getAsIndexOpFoldResult(ctx, sgLayout));
-        if (failed(delinearized))
-          return rewriter.notifyMatchFailure(op, "Failed to delinearize sgId");
-        SmallVector<Value> sgIds = *delinearized;
-
-        SmallVector<int64_t> distUnitShape(sgLayout.size());
-        SmallVector<Value> localOffset(sgLayout.size());
-        for (size_t i = 0; i < sgLayout.size(); i++) {
-          distUnitShape[i] = std::min(sgLayout[i] * sgShape[i], shape[i]);
-          localOffset[i] = rewriter.createOrFold<index::MulOp>(
-              loc, sgIds[i],
-              rewriter.create<arith::ConstantIndexOp>(loc, sgShape[i]));
-        }
-
-        auto tdescTy = xegpu::TensorDescType::get(
-            sgShape, type.getElementType(), 1, false, xegpu::MemorySpace::SLM,
-            target.dropSgLayoutAndData());
-        auto valueTy = VectorType::get(sgShape, type.getElementType());
-
-        SmallVector<OpFoldResult> zeros = getAsIndexOpFoldResult(
-            ctx, SmallVector<int64_t>(sgLayout.size(), 0));
-        for (auto baseOffsets : StaticTileOffsetRange(shape, distUnitShape)) {
-          SmallVector<OpFoldResult> offsets = calculateGlobalOffsets(
-              rewriter, loc, zeros, localOffset, baseOffsets, distUnitShape);
-          auto tdesc = rewriter.create<xegpu::CreateNdDescOp>(
-              loc, tdescTy, slmBuffer, offsets);
-          auto newOp = rewriter.create<xegpu::LoadNdOp>(
-              loc, TypeRange({valueTy}), ValueRange({tdesc}));
-          values.push_back(newOp);
-        }
-      }
-    }
+    // TODO: currently we only support for optimal case, where input and
+    // output has the same sg_layout and sg_data, so SLM is not involved.
+    if (inputSgLayout != targetSgLayout || inputSgData != targetSgData)
+      return failure();
 
     input = input.dropSgLayoutAndData();
     target = target.dropSgLayoutAndData();
 
-    SmallVector<Value> newOps;
-    for (auto src : values) {
-      auto newOp = rewriter.create<xegpu::ConvertLayoutOp>(
-          op.getLoc(), src.getType(), src, input, target);
-      newOps.push_back(newOp);
+    SmallVector<Value> newOps(adaptor.getSource());
+
+    if (input && target) {
+      for (auto [i, src] : llvm::enumerate(adaptor.getSource())) {
+        auto newOp = rewriter.create<xegpu::ConvertLayoutOp>(
+            op.getLoc(), src.getType(), src, input, target);
+        newOps[i] = newOp;
+      }
     }
-    rewriter.replaceOpWithMultiple(op, newOps);
+    rewriter.replaceOpWithMultiple(op, {newOps});
     return success();
   }
 };

mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp

@@ -124,7 +124,7 @@ xegpu::LayoutAttr xegpu::getLayoutAttr(const Value value) {
     Operation *defOp = result.getDefiningOp();
     assert(defOp && "result must have a defining op");
 
-    // For ConvertLayoutOp, the layout is stored in the tensor descriptor
+    // For ConvertLayoutOp, the layout is stored in the targetLayoutAttr
     if (auto convertOp = dyn_cast<xegpu::ConvertLayoutOp>(defOp))
       return convertOp.getTargetLayoutAttr();
 

mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir

@@ -198,4 +198,14 @@ gpu.module @test_round_robin_assignment {
     gpu.return
   }
 
+  gpu.func @convert_layout_optimal(%arg0: memref<32x64xf32>) {
+    %0 = xegpu.create_nd_tdesc %arg0[0, 0] : memref<32x64xf32> -> !xegpu.tensor_desc<32x64xf32, #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>>
+    //CHECK-2: xegpu.load_nd {{.*}}  : !xegpu.tensor_desc<16x16xf32, #xegpu.layout<inst_data = [16, 16]>> -> vector<16x16xf32>
+    //CHECK-2: xegpu.convert_layout {{.*}} <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x16xf32>
+    %1 = xegpu.load_nd %0  : !xegpu.tensor_desc<32x64xf32, #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>> -> vector<32x64xf32>
+    %2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>,
+                                   target_layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16]>}> : vector<32x64xf32>
+    gpu.return
+  }
+
 }