TheRock/compiler/amd-llvm/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp

//===---- XeGPUBlocking.cpp ---- XeGPU Blocking Pass ----------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/XeGPU/Transforms/Passes.h"

#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/DebugLog.h"

namespace mlir {
namespace xegpu {
#define GEN_PASS_DEF_XEGPUBLOCKING
#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
} // namespace xegpu
} // namespace mlir

#define DEBUG_TYPE "xegpu-blocking"

using namespace mlir;

namespace {

// reslove the unrealized conversion cast ops generated when doing SCF
// Structural Type Conversion. It will have two formats, N:1 vector
// cast and 1:N vector cast. vector::insert_strided_slice ops will be
// used for the first case, and vector::extract_strided_slice ops will be
// used for the second case.
static void
resolveUnrealizedConversionCastOp(UnrealizedConversionCastOp castOp) {
  ValueRange inputs = castOp.getInputs();
  ValueRange outputs = castOp.getOutputs();

  auto hasIdenticalVectorTypes = [](ValueRange values) {
    auto types = values.getTypes();
    return llvm::all_of(types, [&](Type type) {
      return isa<VectorType>(type) && type == types.front();
    });
  };

  // We only interest in the case where all inputs and outputs have the
  // identical VectorTypes
  if (!hasIdenticalVectorTypes(inputs) || !hasIdenticalVectorTypes(outputs)) {
    LDBG() << "skip unrealized conversion cast op not emulating pack/unpack.";
    return;
  }

  VectorType outputTy = dyn_cast<VectorType>(outputs[0].getType());
  OpBuilder builder(castOp);
  if (inputs.size() > 1 && outputs.size() == 1) {
    // the castOp is emulating an unpack op
    ArrayRef<int64_t> shape = outputTy.getShape();
    Value result = xegpu::createVectorWithShapeFromValues(
        builder, castOp.getLoc(), inputs, shape);
    castOp->replaceAllUsesWith(ValueRange(result));
    castOp->erase();
  } else if (castOp.getNumResults() > 1 && castOp.getNumOperands() == 1) {
    // the castOp is emulating a pack op
    ArrayRef<int64_t> tileShape = outputTy.getShape();
    SmallVector<Value> results = xegpu::extractVectorsWithShapeFromValue(
        builder, castOp.getLoc(), inputs[0], tileShape);
    castOp->replaceAllUsesWith(results);
    castOp->erase();
  }
}

//===------------------------------------------------------------------------===//
// The XeGPUBlockingPass leverages the unroll patterns for XeGPU and Vector ops
// to partition operations that process large shapes into multiple operations on
// smaller shapes, as specified by the inst_data in the layout attribute. This
// enables each resulting operation to be efficiently mapped to a hardware
// instruction.
//===------------------------------------------------------------------------===//

class XeGPUBlockingPass final
    : public xegpu::impl::XeGPUBlockingBase<XeGPUBlockingPass> {
public:
  void runOnOperation() override;

private:
  // Get the tile shape for a given OpOperand or OpResult by examining the
  // corresponding layout attribute. If layout is not present or is not a
  // subgroup level layout, it returns std::nullopt.
  template <typename T,
            typename = std::enable_if_t<std::is_same_v<T, OpOperand> ||
                                        std::is_same_v<T, OpResult>>>
  std::optional<SmallVector<int64_t>>
  getTileShape(const T &operandOrResult) const;

  // Get the tile shape for a given operation.
  std::optional<SmallVector<int64_t>> getTileShape(Operation *op) const;

  // Determine if the operation requires unrolling. Return false if all operands
  // and results have tile shapes identical to their original types. Otherwise,
  // return true.
  bool needsUnroll(Operation *op) const;
};
} // namespace

template <typename T, typename>
std::optional<SmallVector<int64_t>>
XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
  Value value;
  if constexpr (std::is_same_v<T, OpOperand>) {
    value = operandOrResult.get();
  } else {
    value = (Value)operandOrResult;
  }

  xegpu::DistributeLayoutAttr layout =
      xegpu::getDistributeLayoutAttr(operandOrResult);
  if (layout && layout.isForSubgroup()) {
    if (!layout.getEffectiveInstDataAsInt().empty()) {
      SmallVector<int64_t> instData = layout.getEffectiveInstDataAsInt();
      return instData;
    }
    if (auto type = dyn_cast<ShapedType>(value.getType()))
      return llvm::to_vector(type.getShape());
  }
  LDBG() << "failed to getTileShape for: " << value;
  return std::nullopt;
}

std::optional<SmallVector<int64_t>>
XeGPUBlockingPass::getTileShape(Operation *op) const {
  if (isa<xegpu::CreateNdDescOp, xegpu::LoadMatrixOp>(op))
    return getTileShape(op->getOpResult(0));
  if (isa<xegpu::PrefetchNdOp, xegpu::LoadNdOp, xegpu::PrefetchOp,
          xegpu::StoreMatrixOp>(op))
    return getTileShape(op->getOpOperand(0));
  if (isa<xegpu::StoreNdOp>(op))
    return getTileShape(op->getOpOperand(1));

  if (isa<xegpu::LoadGatherOp>(op))
    return getTileShape(op->getOpResult(0));

  if (auto convertLayoutOp = dyn_cast<xegpu::ConvertLayoutOp>(op)) {
    auto inputInstData =
        convertLayoutOp.getInputLayout().getEffectiveInstDataAsInt();
    auto targetInstData =
        convertLayoutOp.getTargetLayout().getEffectiveInstDataAsInt();
    // return the one with larger size
    if (computeProduct(inputInstData) >= computeProduct(targetInstData))
      return inputInstData;
    else
      return targetInstData;
  }

  if (isa<xegpu::StoreScatterOp>(op))
    return getTileShape(op->getOpOperand(0));

  // Helper lambda to validate and get A/B tiles
  auto validateABTiles = [&](Operation *op)
      -> std::optional<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>> {
    std::optional<SmallVector<int64_t>> aTile =
        getTileShape(op->getOpOperand(0));
    std::optional<SmallVector<int64_t>> bTile =
        getTileShape(op->getOpOperand(1));

    if (!aTile || aTile->size() != 2 || !bTile || bTile->size() != 2)
      return std::nullopt;

    // semantic check for A and B
    if ((*aTile)[1] != (*bTile)[0])
      return std::nullopt;

    return std::make_pair(*aTile, *bTile);
  };

  // Helper lambda to validate C tile
  auto validateCTile = [&](Operation *op, unsigned cOperandIdx,
                           const SmallVector<int64_t> &aTile,
                           const SmallVector<int64_t> &bTile) -> bool {
    if (op->getNumOperands() <= cOperandIdx)
      return true;

    std::optional<SmallVector<int64_t>> cTile =
        getTileShape(op->getOpOperand(cOperandIdx));
    int64_t expectedCTile[2] = {aTile[0], bTile[1]};
    if (!cTile || !llvm::equal(*cTile, expectedCTile))
      return false;
    return true;
  };

  // Helper lambda to validate scale A tile for DpasMxOp
  auto validateScaleATile =
      [&](Operation *op, unsigned scaleAOperandIdx,
          const SmallVector<int64_t> &aTile) -> std::optional<int64_t> {
    std::optional<SmallVector<int64_t>> aScaleTile =
        getTileShape(op->getOpOperand(scaleAOperandIdx));

    if (!aScaleTile || aScaleTile->size() != 2)
      return std::nullopt;

    // Validate scale_a tile: [M_tile, K_scale]
    // M dimension must match A's M dimension
    if ((*aScaleTile)[0] != aTile[0])
      return std::nullopt;

    // Return the K scale factor
    return (*aScaleTile)[1];
  };

  // Helper lambda to validate scale B tile for DpasMxOp
  auto validateScaleBTile =
      [&](Operation *op, unsigned scaleBOperandIdx,
          const SmallVector<int64_t> &bTile) -> std::optional<int64_t> {
    std::optional<SmallVector<int64_t>> bScaleTile =
        getTileShape(op->getOpOperand(scaleBOperandIdx));

    if (!bScaleTile || bScaleTile->size() != 2)
      return std::nullopt;

    // Validate scale_b tile: [K_scale, N_tile]
    // N dimension must match B's N dimension
    if ((*bScaleTile)[1] != bTile[1])
      return std::nullopt;

    // Return the K scale factor
    return (*bScaleTile)[0];
  };

  if (isa<xegpu::DpasOp>(op)) {
    auto abTiles = validateABTiles(op);
    if (!abTiles)
      return std::nullopt;

    auto [aTile, bTile] = *abTiles;

    // semantic check for C
    if (!validateCTile(op, 2, aTile, bTile))
      return std::nullopt;

    return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1]});
  }

  if (auto dpasMxOp = dyn_cast<xegpu::DpasMxOp>(op)) {
    auto abTiles = validateABTiles(op);
    if (!abTiles)
      return std::nullopt;

    auto [aTile, bTile] = *abTiles;

    // Validate C tile if present using op-specific accessor
    if (dpasMxOp.getAcc()) {
      unsigned accOperandIdx = 2; // acc is the 3rd operand
      if (!validateCTile(op, accOperandIdx, aTile, bTile))
        return std::nullopt;
    }

    // Validate scale tiles if present using op-specific accessors
    int64_t kScaleFactor = 1;
    std::optional<int64_t> scaleAFactor;
    std::optional<int64_t> scaleBFactor;

    if (dpasMxOp.getScaleA()) {
      unsigned scaleAOperandIdx = 2 + (dpasMxOp.getAcc() ? 1 : 0);
      scaleAFactor = validateScaleATile(op, scaleAOperandIdx, aTile);
      if (!scaleAFactor)
        return std::nullopt;
    }

    if (dpasMxOp.getScaleB()) {
      unsigned scaleBOperandIdx =
          2 + (dpasMxOp.getAcc() ? 1 : 0) + (dpasMxOp.getScaleA() ? 1 : 0);
      scaleBFactor = validateScaleBTile(op, scaleBOperandIdx, bTile);
      if (!scaleBFactor)
        return std::nullopt;
    }

    // If both scales are present, their K dimensions must match
    if (scaleAFactor && scaleBFactor) {
      if (*scaleAFactor != *scaleBFactor)
        return std::nullopt;
      kScaleFactor = *scaleAFactor;
    } else if (scaleAFactor) {
      kScaleFactor = *scaleAFactor;
    } else if (scaleBFactor) {
      kScaleFactor = *scaleBFactor;
    }

    return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1], kScaleFactor});
  }

  if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)
    return getTileShape(op->getOpResult(0));

  if (isa<vector::MultiDimReductionOp>(op))
    return getTileShape(op->getOpOperand(0));

  if (isa<vector::TransposeOp, vector::BroadcastOp, vector::StepOp,
          vector::ShapeCastOp, vector::ConstantMaskOp, vector::CreateMaskOp,
          vector::BitCastOp, vector::InterleaveOp, vector::DeinterleaveOp>(op))
    return getTileShape(op->getOpResult(0));

  return std::nullopt;
}

bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
  // skip the op if any of its operands or results has workgroup level layouts
  bool hasWgLayoutOperands =
      llvm::any_of(op->getOpOperands(), [](OpOperand &opr) {
        xegpu::DistributeLayoutAttr layout =
            xegpu::getDistributeLayoutAttr(opr);
        return layout && layout.isForWorkgroup();
      });
  bool hasWgLayoutResults =
      llvm::any_of(op->getOpResults(), [](OpResult result) {
        xegpu::DistributeLayoutAttr layout =
            xegpu::getDistributeLayoutAttr(result);
        return layout && layout.isForWorkgroup();
      });
  if (hasWgLayoutOperands || hasWgLayoutResults) {
    LDBG() << "skip unrolling for op with workgroup level layout: " << *op;
    return false;
  }

  auto isUnrollable = [](Value value, ArrayRef<int64_t> tileShape) {
    Type valTy = value.getType();
    if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(valTy)) {
      xegpu::DistributeLayoutAttr layout = tdescTy.getLayoutAttr();
      return layout && !layout.getEffectiveInstDataAsInt().empty();
    }
    auto shapedType = dyn_cast<ShapedType>(valTy);
    return shapedType && !llvm::equal(tileShape, shapedType.getShape());
  };

  bool hasUnrollableOperands =
      llvm::any_of(op->getOpOperands(), [&](OpOperand &opr) {
        std::optional<SmallVector<int64_t>> tileShape = getTileShape(opr);
        return tileShape.has_value() && isUnrollable(opr.get(), *tileShape);
      });
  bool hasUnrollableResults =
      llvm::any_of(op->getOpResults(), [&](OpResult result) {
        std::optional<SmallVector<int64_t>> tileShape = getTileShape(result);
        return tileShape.has_value() && isUnrollable(result, *tileShape);
      });
  // ConvertLayoutOp must be processed to drop the inst_data in the layout
  bool isConvertLayoutWithInstData = false;
  if (auto convertLayoutOp = dyn_cast<xegpu::ConvertLayoutOp>(op)) {
    auto targettLayout = convertLayoutOp.getTargetLayout();
    if (targettLayout && !targettLayout.getEffectiveInstDataAsInt().empty()) {
      isConvertLayoutWithInstData = true;
    }
  }
  return hasUnrollableOperands || hasUnrollableResults ||
         isConvertLayoutWithInstData;
}

void XeGPUBlockingPass::runOnOperation() {
  MLIRContext *ctx = &getContext();
  Operation *op = getOperation();

  if (!xegpu::recoverTemporaryLayouts(op)) {
    signalPassFailure();
    return;
  }

  auto getTileShapeAndCount = [](llvm::ArrayRef<int64_t> shape,
                                 xegpu::DistributeLayoutAttr layout) {
    int count = 1;
    SmallVector<int64_t> tileShape(shape);
    if (layout && !layout.getEffectiveInstDataAsInt().empty()) {
      tileShape = layout.getEffectiveInstDataAsInt();
      count = computeProduct(shape) / computeProduct(tileShape);
    }
    return std::make_pair(tileShape, count);
  };

  // Perform type conversion for SCF control folow ops
  TypeConverter converter;
  converter.addConversion([](Type type) -> Type { return type; });
  converter.addConversion(
      [&](RankedTensorType type,
          SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
        Type elemTy = type.getElementType();
        ArrayRef<int64_t> shape = type.getShape();

        auto layout =
            llvm::dyn_cast_if_present<xegpu::LayoutAttr>(type.getEncoding());
        if (layout && layout.isForWorkgroup())
          return failure();

        int count;
        SmallVector<int64_t> subShape;
        std::tie(subShape, count) = getTileShapeAndCount(shape, layout);
        auto newTy = VectorType::get(subShape, elemTy);
        result.append(count, newTy);
        return success();
      });
  converter.addConversion(
      [&](xegpu::TensorDescType type,
          SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
        Type elemTy = type.getElementType();
        ArrayRef<int64_t> shape = type.getShape();

        xegpu::DistributeLayoutAttr layout = type.getLayoutAttr();
        if (layout && layout.isForWorkgroup())
          return failure();

        int count;
        SmallVector<int64_t> subShape;
        std::tie(subShape, count) = getTileShapeAndCount(shape, layout);

        if (layout)
          layout = layout.dropInstData();

        auto newTy = xegpu::TensorDescType::get(
            type.getContext(), subShape, elemTy, type.getEncoding(), layout);
        result.append(count, newTy);
        return success();
      });

  xegpu::doSCFStructuralTypeConversionWithTensorType(op, converter);

  xegpu::UnrollOptions options;
  options.setFilterConstraint(
      [&](Operation *op) -> LogicalResult { return success(needsUnroll(op)); });

  options.setNativeShapeFn([&](Operation *op) { return getTileShape(op); });

  options.setUnrolledTypesFn([&](ShapedType type, ArrayRef<int64_t> tileShape,
                                 bool returnSingleType = false) {
    Type elemTy = type.getElementType();
    Type newTy;

    if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(type)) {

      Attribute encoding = tdescTy.getEncoding();

      newTy =
          xegpu::TensorDescType::get(ctx, tileShape, elemTy, encoding,
                                     tdescTy.getLayoutAttr().dropInstData());
    } else {
      newTy = VectorType::get(tileShape, elemTy);
    }

    if (returnSingleType)
      return SmallVector<Type>{newTy};
    std::optional<SmallVector<int64_t>> ratio =
        computeShapeRatio(type.getShape(), tileShape);
    assert(ratio && "The shape of the type must be a multiple of tileShape.");
    return SmallVector<Type>(computeProduct(*ratio), newTy);
  });

  RewritePatternSet patterns(ctx);
  vector::UnrollVectorOptions vectorOptions;
  vectorOptions.setNativeShapeFn(options.nativeShape);

  populateXeGPUUnrollPatterns(patterns, options);
  vector::populateVectorUnrollPatterns(patterns, vectorOptions);

  // Note: The pattern driver does op folding as well and clean up.
  // But intermediate insert/extract strided slice ops with
  // unrealized conversion cast ops in the middle does not get
  // cleaned up in this step. One more round of folding is needed
  // after the walk to resolve those unrealized conversion cast ops.
  (void)applyPatternsGreedily(op, std::move(patterns));

  op->walk([](Operation *op) {
    // Remove the layout attributes cached per operands.
    for (OpOperand &opr : op->getOpOperands()) {
      std::string name = xegpu::getTemporaryLayoutName(opr);
      if (op->hasAttrOfType<xegpu::DistributeLayoutAttr>(name))
        op->removeAttr(name);
    }

    // Update the layout attributes per result.
    for (OpResult result : op->getOpResults()) {
      std::string name = xegpu::getTemporaryLayoutName(result);
      if (auto layout = op->getAttrOfType<xegpu::DistributeLayoutAttr>(name)) {
        op->removeAttr(name);
        if (!isa<LoopLikeOpInterface>(op))
          xegpu::setDistributeLayoutAttr(result, layout.dropInstData());
      }
    }

    // Drop left-over inst_data if the unroll pattern does not being applied,
    // say, inst_data just matches their shape.
    SmallVector<NamedAttribute> newAttrs =
        xegpu::dropInstDataOnAttrs(op->getAttrs());
    op->setAttrs(newAttrs);

    // Resolve unrealized conversion cast ops emulating pack/unpack
    if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op))
      resolveUnrealizedConversionCastOp(castOp);
  });

  // One more round of folding to clean up the intermediate
  // insert/extract strided slice ops.
  RewritePatternSet emptyPatterns(ctx);
  (void)applyPatternsGreedily(op, std::move(emptyPatterns));
}