XeGPUOps.cpp

Go to the documentation of this file.

//===- XeGPUOps.cpp - MLIR XeGPU ops implementation -------------*- C++ -*-===//
//
// 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/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
 
#include "llvm/Support/Debug.h"
 
#include <utility>
 
#define DEBUG_TYPE "xegpu"
 
using namespace mlir;
using namespace mlir::xegpu;
 
template <typename T>
static std::string makeString(const T &array, bool breakline = false) {
  std::string buf;
  buf.clear();
  llvm::raw_string_ostream os(buf);
  os << "[";
  for (size_t i = 1; i < array.size(); i++) {
    os << array[i - 1] << ", ";
    if (breakline)
      os << "\n\t\t";
  }
  os << array.back() << "]";
  return buf;
}
 
static SmallVector<int64_t> getShapeOf(Type type) {
  SmallVector<int64_t> shape;
  if (auto ty = llvm::dyn_cast<ShapedType>(type))
    shape = SmallVector<int64_t>(ty.getShape());
  else
    shape.push_back(1);
  return shape;
}
 
static bool isReadHintOrNone(const CachePolicyAttr &attr) {
  if (!attr)
    return true;
  auto kind = attr.getValue();
  return kind == CachePolicy::CACHED || kind == CachePolicy::UNCACHED ||
         kind == CachePolicy::STREAMING || kind == CachePolicy::READ_INVALIDATE;
}
 
static bool isWriteHintOrNone(const CachePolicyAttr &attr) {
  if (!attr)
    return true;
  auto kind = attr.getValue();
  return kind == CachePolicy::CACHED || kind == CachePolicy::UNCACHED ||
         kind == CachePolicy::WRITE_BACK || kind == CachePolicy::WRITE_THROUGH;
}
 
static LogicalResult
isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
                                 VectorType valueTy,
                                 function_ref<InFlightDiagnostic()> emitError) {
 
  auto maskVecTy = dyn_cast<VectorType>(maskTy);
  auto offsetsVecTy = dyn_cast<VectorType>(offsetsTy);
 
  // The offsets and mask describe the same lane grid: either both are scalar
  // (single lane) or both are vectors carrying one offset and one mask bit per
  // lane, so they must have the same shape.
  if (static_cast<bool>(maskVecTy) != static_cast<bool>(offsetsVecTy))
    return emitError() << "Expecting offsets and mask to both be scalar or "
                          "both be vectors.";
  if (maskVecTy && getShapeOf(maskTy) != getShapeOf(offsetsTy))
    return emitError() << "Expecting offsets and mask to have the same shape.";
 
  // Scalar payload (e.g. `index, i1 -> f16`): the offsets and mask must be
  // scalar as well.
  if (!valueTy) {
    if (maskVecTy || offsetsVecTy)
      return emitError() << "Expecting scalar mask and offsets.";
    return success();
  }
 
  // One offset and one mask bit per element, so the value has the same shape as
  // the mask and the offsets. A scalar and a size-1 vector both stand for a
  // single element, so they are interchangeable.
  int64_t maskSize = maskVecTy ? maskVecTy.getNumElements() : 1;
  if (valueTy.getNumElements() != maskSize ||
      (maskVecTy && getShapeOf(valueTy) != getShapeOf(maskTy)))
    return emitError() << "Value shape must match mask shape.";
  return success();
}
 
// Validates the `contiguity` attribute against the op's offsets type: the
// innermost offsets dimension is contiguous in runs of `size`, so `size` must
// be >= 2 and must divide that dimension.
static LogicalResult
isValidContiguity(std::optional<uint64_t> contiguity, Type offsetsTy,
                  function_ref<InFlightDiagnostic()> emitError) {
  if (!contiguity)
    return success();
  auto offsetsVecTy = dyn_cast<VectorType>(offsetsTy);
  if (!offsetsVecTy)
    return emitError() << "contiguity requires vector offsets (one per lane).";
  int64_t size = static_cast<int64_t>(*contiguity);
  int64_t inner = offsetsVecTy.getShape().back();
  if (size < 2)
    return emitError() << "contiguity = " << size << " (must be >= 2)";
  if (inner % size != 0)
    return emitError() << "contiguity = " << size
                       << " (must divide the innermost offsets dim " << inner
                       << ")";
  return success();
}
 
LogicalResult
IsValidMatrixOpParams(VectorType dataTy, MemDescType mdescTy,
                      UnitAttr subgroup_block_io, DistributeLayoutAttr layout,
                      function_ref<InFlightDiagnostic()> emitError) {
 
  if (!dataTy) {
    if (subgroup_block_io)
      return emitError() << "subgroup_block_io "
                            "are only allowed when result is a VectorType.";
    else
      return success();
  }
 
  ArrayRef<int64_t> dataShape = dataTy.getShape();
  ArrayRef<int64_t> mdescShape = mdescTy.getShape();
 
  SmallVector<int64_t> blockShape = mdescTy.getBlockShape();
  ArrayAttr strideAttr = mdescTy.getStrideAttr();
  SmallVector<int64_t> strides;
  for (Attribute attr : strideAttr.getValue()) {
    strides.push_back(cast<IntegerAttr>(attr).getInt());
  }
  if (subgroup_block_io && layout) {
    auto laneData = layout.getEffectiveLaneDataAsInt();
    auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
    if (!laneData.empty()) {
      bool isLaneDataContiguous =
          std::all_of(laneData.begin(), std::prev(laneData.end()),
                      [](int x) { return x == 1; });
      if (!isLaneDataContiguous)
        return emitError() << "With subgroup_block_io, accessed data must be "
                              "contiguous and coalesced.";
      for (size_t i = 0; i < laneData.size(); ++i) {
        if (laneLayout[i] != blockShape[i])
          return emitError() << "With subgroup_block_io, the block shape must "
                                "match the lane layout.";
        if (laneLayout[i] != 1 && strides[i] != 1)
          return emitError() << "With subgroup_block_io, the distributed "
                                "dimensions must be contiguous.";
      }
    }
  }
 
  if (layout && !layout.isDistributable(
                    SmallVector<int64_t>(dataShape.begin(), dataShape.end())))
    return emitError() << "Value shape is not distributable with the layout";
 
  if (dataShape.size() == mdescShape.size()) {
    if (llvm::any_of(llvm::zip_equal(dataShape, mdescShape),
                     [](auto p) { return std::get<0>(p) > std::get<1>(p); }))
      return emitError() << "data shape must not exceed mem_desc shape.";
  }
  // if the subgroup_block_io attribute is set, mdescTy must have block
  // attribute
  if (subgroup_block_io && !blockShape.size())
    return emitError() << "mem_desc must have block attribute when "
                          "subgroup_block_io is set.";
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_CreateMemDescOp
//===----------------------------------------------------------------------===//
LogicalResult CreateMemDescOp::verify() {
  auto srcTy = getSource().getType();
  if (!memref::isStaticShapeAndContiguousRowMajor(srcTy))
    return emitOpError("source memref must be contiguous.");
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_CreateNdDescOp
//===----------------------------------------------------------------------===//
 
void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
                           Type tdesc, TypedValue<MemRefType> source) {
  build(builder, state, tdesc, source, ValueRange({}) /* empty dynamic shape */,
        ValueRange({}) /* empty dynamic strides */,
        DenseI64ArrayAttr({}) /* empty const shape*/,
        DenseI64ArrayAttr({}) /* empty const strides*/);
}
 
void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
                           Type tdesc, Value source,
                           llvm::ArrayRef<OpFoldResult> shape,
                           llvm::ArrayRef<OpFoldResult> strides) {
  Type srcTy = source.getType();
  assert((isa<IntegerType, MemRefType>(srcTy)) &&
         "Source has to be either int or memref.");
 
  llvm::SmallVector<Value> dynamicShape;
  llvm::SmallVector<Value> dynamicStrides;
 
  llvm::SmallVector<int64_t> staticShape;
  llvm::SmallVector<int64_t> staticStrides;
 
  dispatchIndexOpFoldResults(shape, dynamicShape, staticShape);
  dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides);
 
  auto staticShapeAttr = builder.getDenseI64ArrayAttr(staticShape);
  auto staticStridesAttr = builder.getDenseI64ArrayAttr(staticStrides);
 
  if (auto memrefTy = dyn_cast<MemRefType>(srcTy)) {
    auto memrefShape = memrefTy.getShape();
    auto [memrefStrides, _] = memrefTy.getStridesAndOffset();
 
    // if shape and strides are from Memref, we don't need attributes for them
    // to keep the IR print clean (only do so for full-static case, otherwise
    // printer would fail trying to print empty array-attr).
    if (staticShape == memrefShape && staticStrides == memrefStrides &&
        dynamicShape.empty() && dynamicStrides.empty()) {
      staticShapeAttr = DenseI64ArrayAttr();
      staticStridesAttr = DenseI64ArrayAttr();
    }
  }
 
  build(builder, state, tdesc, source, dynamicShape, dynamicStrides,
        staticShapeAttr, staticStridesAttr);
}
 
LogicalResult CreateNdDescOp::verify() {
  auto srcMemrefTy = dyn_cast<MemRefType>(getSourceType());
  size_t rank = srcMemrefTy ? srcMemrefTy.getRank() : getMixedSizes().size();
  bool invalidElemTy = false;
 
  // Memory space of created TensorDesc should match with the source.
  // Both source and TensorDesc are considered for global memory by default,
  // if the memory scope attr is not specified. If source is an integer,
  // it is considered as ptr to global memory.
  auto srcMemorySpace = getSourceMemorySpace();
  auto tdescMemorySpace = static_cast<unsigned>(getType().getMemorySpace());
  if (srcMemorySpace != tdescMemorySpace)
    return emitOpError("Memory space mismatch.")
           << " Source: " << srcMemorySpace
           << ", TensorDesc: " << tdescMemorySpace;
 
  // check source type matches the rank if it is a memref.
  // It also should have the same ElementType as TensorDesc.
  if (auto memrefTy = dyn_cast<MemRefType>(getSourceType()))
    invalidElemTy |= memrefTy.getElementType() != getElementType();
 
  bool hasExplicitShapeStrides =
      !getShape().empty() || !getStrides().empty() ||
      (getConstShapeAttr() && !getConstShapeAttr().empty()) ||
      (getConstStridesAttr() && !getConstStridesAttr().empty());
 
  if (llvm::isa<IntegerType>(getSourceType())) {
    // strides and shape must present for integer source.
    if (getMixedStrides().empty() || getMixedSizes().empty())
      return emitOpError("expecting strides and shape to be present for "
                         "integer source.");
    if (getMixedSizes().size() != getMixedStrides().size())
      return emitOpError("Expecting the rank of shape and strides to match.");
  } else if (srcMemrefTy && hasExplicitShapeStrides) {
    return emitOpError("shape and strides should not be specified for a memref "
                       "source; they are inferred from the memref.");
  }
 
  // check result TensorDesc rank
  if (getType().getRank() > (int64_t)rank)
    return emitOpError("Expecting the TensorDesc rank is not greater than the "
                       "ranks of shape, strides or the memref source.");
 
  if (invalidElemTy)
    return emitOpError("TensorDesc should have the same element "
                       "type with the source if it is a memref.\n");
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_PrefetchNdOp
//===----------------------------------------------------------------------===//
 
void PrefetchNdOp::build(OpBuilder &builder, OperationState &state,
                         Value tensorDesc, ArrayRef<OpFoldResult> offsets,
                         xegpu::CachePolicyAttr l1_hint,
                         xegpu::CachePolicyAttr l2_hint,
                         xegpu::CachePolicyAttr l3_hint,
                         xegpu::DistributeLayoutAttr layout) {
  SmallVector<Value> dynamicOffsets;
  SmallVector<int64_t> staticOffsets;
  dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
 
  auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);
 
  build(builder, state, tensorDesc, dynamicOffsets, staticOffsetsAttr, l1_hint,
        l2_hint, l3_hint, /*anchor_layout=*/layout);
}
 
LogicalResult PrefetchNdOp::verify() {
  auto tdescTy = getTensorDescType();
 
  if (!isReadHintOrNone(getL1HintAttr()))
    return emitOpError("invalid l1_hint: ") << getL1HintAttr();
 
  if (!isReadHintOrNone(getL2HintAttr()))
    return emitOpError("invalid l2_hint: ") << getL2HintAttr();
 
  if (!isReadHintOrNone(getL3HintAttr()))
    return emitOpError("invalid l3_hint: ") << getL3HintAttr();
 
  int64_t tDescRank = tdescTy.getRank();
  int64_t offsetSize = getMixedOffsets().size();
  if (offsetSize != tDescRank)
    return emitOpError(
        "Mismatched ranks between offsets and tensor descriptor");
 
  if (auto layout = getAnchorLayout()) {
    if (!layout.isDistributable(getShapeOf(tdescTy)))
      return emitOpError(
          "TensorDesc shape is not distributable with the layout");
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_LoadNdOp
//===----------------------------------------------------------------------===//
 
void LoadNdOp::build(OpBuilder &builder, OperationState &state, Type retType,
                     Value tensorDesc, ArrayRef<OpFoldResult> offsets,
                     UnitAttr packed, DenseI64ArrayAttr transpose,
                     xegpu::CachePolicyAttr l1_hint,
                     xegpu::CachePolicyAttr l2_hint,
                     xegpu::CachePolicyAttr l3_hint,
                     xegpu::DistributeLayoutAttr layout) {
  SmallVector<Value> dynamicOffsets;
  SmallVector<int64_t> staticOffsets;
  dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
 
  auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);
 
  build(builder, state, retType, tensorDesc, dynamicOffsets, staticOffsetsAttr,
        packed, transpose, l1_hint, l2_hint, l3_hint,
        /*anchor_layout=*/layout);
}
 
LogicalResult LoadNdOp::verify() {
  auto tdescTy = getTensorDescType();
  auto valueTy = getType();
 
  if (!valueTy)
    return emitOpError("Invalid result, it should be a VectorType.\n");
 
  if (!isReadHintOrNone(getL1HintAttr()))
    return emitOpError("invalid l1_hint: ") << getL1HintAttr();
 
  if (!isReadHintOrNone(getL2HintAttr()))
    return emitOpError("invalid l2_hint: ") << getL2HintAttr();
 
  if (!isReadHintOrNone(getL3HintAttr()))
    return emitOpError("invalid l3_hint: ") << getL3HintAttr();
 
  int tdescElems = tdescTy.getNumElements() * tdescTy.getArrayLength();
  int valueElems = valueTy.getNumElements();
 
  // If the result vector is 1D and has less elements than the tensor
  // descriptor, it is supposed to be a SIMT op. The layout attribute in
  // tensor_desc is not needed.
  if (valueElems < tdescElems && valueTy.getRank() == 1) {
    // SIMT mode doesn't need LayoutAttr.
    if (tdescTy.getLayoutAttr())
      return emitOpError()
             << "TensorDesc doesn't need LayoutAttr for SIMT code";
 
    // For SIMT code, the load is evenly distributed across all lanes in a
    // subgroup. Since subgroup size is arch dependent, we only check even
    // distribution here.
    if (tdescElems % valueElems)
      return emitOpError()
             << "Result shape " << makeString(getShapeOf(valueTy))
             << " is not a valid distribution for tensor descriptor "
             << tdescTy;
 
    return success();
  }
 
  // Check SIMD mode.
  auto tdescShape = getShapeOf(tdescTy);
  auto valueShape = getShapeOf(valueTy);
 
  if (getTranspose()) {
    auto trans = getTranspose().value();
    // Make sure the transpose value is valid, and apply it
    if (llvm::all_of(trans, [&](size_t s) { return s < tdescShape.size(); }))
      tdescShape = applyPermutation(tdescShape, trans);
    else
      mlir::emitWarning(getLoc()) << "Invalid transpose attr. It is ignored.";
  }
 
  if (getPacked()) {
    if (tdescTy.getRank() == 2) {
      const int axis = 0;
      auto vnni_factor = valueShape.back();
      tdescShape[axis] /= vnni_factor;
      tdescShape.push_back(vnni_factor);
    } else {
      mlir::emitWarning(getLoc())
          << "Invalid Packed Attr. It is ignored (available for 2D "
             "TensorDesc only).";
    }
  }
 
  // Handle array_length. Two result shape conventions are accepted:
  //   * 3D shape: leading array_length dimension prepended, e.g. descriptor
  //     16x16 with array_length=2 -> [2, 16, 16].
  //   * Stacked 2D shape: array blocks stacked along the non-FCD (first)
  //     dimension, e.g. descriptor 16x16 with array_length=2 -> [32, 16].
  auto array_len = tdescTy.getArrayLength();
  SmallVector<int64_t> stacked2DShape(tdescShape);
  SmallVector<int64_t> threeDShape(tdescShape);
  if (array_len > 1 && !tdescShape.empty()) {
    stacked2DShape[0] *= array_len;
    threeDShape.insert(threeDShape.begin(), array_len);
  }
 
  if (valueShape != stacked2DShape && valueShape != threeDShape)
    return emitOpError() << "Result shape " << makeString(valueShape)
                         << " is not consistent with tensor descriptor "
                         << tdescTy;
 
  int64_t tDescRank = tdescTy.getRank();
  int64_t offsetSize = getMixedOffsets().size();
  if (offsetSize != tDescRank)
    return emitOpError(
        "Mismatched ranks between offsets and tensor descriptor");
 
  if (auto layout = getAnchorLayout()) {
    if (!layout.isDistributable(getShapeOf(tdescTy)))
      return emitOpError(
          "TensorDesc shape is not distributable with the layout");
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_StoreNdOp
//===----------------------------------------------------------------------===//
 
void StoreNdOp::build(OpBuilder &builder, OperationState &state, Value value,
                      Value tensorDesc, ArrayRef<OpFoldResult> offsets,
                      xegpu::CachePolicyAttr l1_hint,
                      xegpu::CachePolicyAttr l2_hint,
                      xegpu::CachePolicyAttr l3_hint,
                      xegpu::DistributeLayoutAttr layout) {
  SmallVector<Value> dynamicOffsets;
  SmallVector<int64_t> staticOffsets;
  dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
 
  auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);
 
  build(builder, state, value, tensorDesc, dynamicOffsets, staticOffsetsAttr,
        l1_hint, l2_hint, l3_hint, /*anchor_layout=*/layout);
}
 
LogicalResult StoreNdOp::verify() {
  auto dstTy = getTensorDescType(); // Tile
  auto valTy = getValueType();      // Vector
 
  if (!valTy)
    return emitOpError("Expecting a VectorType result.\n");
 
  if (!isWriteHintOrNone(getL1HintAttr()))
    return emitOpError("invalid l1_hint: ") << getL1HintAttr();
 
  if (!isWriteHintOrNone(getL2HintAttr()))
    return emitOpError("invalid l2_hint: ") << getL2HintAttr();
 
  if (!isWriteHintOrNone(getL3HintAttr()))
    return emitOpError("invalid l3_hint: ") << getL3HintAttr();
 
  auto array_len = dstTy.getArrayLength();
  if (array_len > 1)
    return emitOpError("array length is not supported by store_nd.\n");
 
  auto tdescElems = dstTy.getNumElements();
  auto valueElems = valTy.getNumElements();
 
  // Similar to LoadNdOp, if the value vector is 1D and has less elements than
  // the tensor descriptor, it is supposed to be a SIMT op. The layout attribute
  // in tensor_desc is not needed.
  if (valTy.getRank() == 1 && valueElems < tdescElems) {
    // SIMT mode doesn't need LayoutAttr.
    if (dstTy.getLayoutAttr())
      return emitOpError()
             << "TensorDesc doesn't need LayoutAttr for SIMT code";
 
    if (tdescElems % valueElems)
      return emitOpError()
             << "Value shape " << makeString(getShapeOf(valTy))
             << " is not a valid distribution for tensor descriptor " << dstTy;
 
    return success();
  }
 
  // SIMD code should have the same shape as the tensor descriptor.
  auto tdescShape = getShapeOf(dstTy);
  auto valueShape = getShapeOf(valTy);
  if (tdescShape != valueShape)
    return emitOpError() << "Value shape " << makeString(valueShape)
                         << " is not consistent with tensor descriptor "
                         << dstTy;
 
  int64_t tDescRank = dstTy.getRank();
  int64_t offsetSize = getMixedOffsets().size();
  if (offsetSize != tDescRank)
    return emitOpError(
        "Mismatched ranks between offsets and tensor descriptor");
 
  if (auto layout = getAnchorLayout()) {
    if (!layout.isDistributable(std::move(tdescShape)))
      return emitOpError(
          "TensorDesc shape is not distributable with the layout");
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_PrefetchOp
//===----------------------------------------------------------------------===//
LogicalResult PrefetchOp::verify() {
  if (!isReadHintOrNone(getL1HintAttr()))
    return emitOpError("invalid l1_hint: ") << getL1HintAttr();
 
  if (!isReadHintOrNone(getL2HintAttr()))
    return emitOpError("invalid l2_hint: ") << getL2HintAttr();
 
  if (!isReadHintOrNone(getL3HintAttr()))
    return emitOpError("invalid l3_hint: ") << getL3HintAttr();
 
  auto srcTy = getSourceType();
  if (srcTy.isInteger() && !getOffsetAlignByteAttr())
    return emitOpError("offset_align_byte is required with integer source.");
 
  if (getOffsetAlignByteAttr() && !srcTy.isInteger())
    return emitOpError("offset_align_byte only allowed with integer source.");
 
  if (auto layout = getAnchorLayout()) {
    // get the offset operand and its shape
    auto offsetsTy = getOffsets().getType();
    if (llvm::isa<VectorType>(offsetsTy) &&
        !layout.isDistributable(getShapeOf(offsetsTy)))
      return emitOpError("offset shape is not distributable with the layout");
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_LoadGatherOp
//===----------------------------------------------------------------------===//
LogicalResult LoadGatherOp::verify() {
  auto maskTy = getMaskType();
  auto valueTy = getValueType();
 
  if (!isReadHintOrNone(getL1HintAttr()))
    return emitOpError("invalid l1_hint: ") << getL1HintAttr();
 
  if (!isReadHintOrNone(getL2HintAttr()))
    return emitOpError("invalid l2_hint: ") << getL2HintAttr();
 
  if (!isReadHintOrNone(getL3HintAttr()))
    return emitOpError("invalid l3_hint: ") << getL3HintAttr();
 
  auto srcTy = getSourceType();
  auto memTy = dyn_cast<MemRefType>(srcTy);
 
  if (memTy && (getElementType() != memTy.getElementType()))
    return emitError() << "Value should have the same element type as MemRef.";
 
  if (auto layout = getAnchorLayout()) {
    if (!layout.isDistributable(getShapeOf(valueTy)))
      return emitOpError("Value shape is not distributable with the layout");
  }
 
  auto offsetsTy = getOffsets().getType();
  if (failed(isValidContiguity(getContiguity(), offsetsTy,
                               [&]() { return emitOpError(); })))
    return failure();
  return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy,
                                          [&]() { return emitOpError(); });
}
 
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
                         Type valueType, Value source,
                         ArrayRef<OpFoldResult> offsets, Value mask,
                         xegpu::CachePolicyAttr l1_hint,
                         xegpu::CachePolicyAttr l2_hint,
                         xegpu::CachePolicyAttr l3_hint) {
  auto loc = source.getLoc();
  int64_t size = static_cast<int64_t>(offsets.size());
  auto type = VectorType::get(size, builder.getIndexType());
  auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
  auto offset = vector::FromElementsOp::create(builder, loc, type, values);
 
  build(builder, state, valueType, source, offset, mask, l1_hint, l2_hint,
        l3_hint, /*anchor_layout=*/nullptr,
        /*contiguity=*/nullptr);
}
 
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
                         Type valueType, Value source,
                         ArrayRef<OpFoldResult> offsets, Value mask,
                         xegpu::CachePolicyAttr l1_hint,
                         xegpu::CachePolicyAttr l2_hint,
                         xegpu::CachePolicyAttr l3_hint,
                         DistributeLayoutAttr layout) {
  auto loc = source.getLoc();
  int64_t size = static_cast<int64_t>(offsets.size());
  auto type = VectorType::get(size, builder.getIndexType());
  auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
  auto offset = vector::FromElementsOp::create(builder, loc, type, values);
 
  build(builder, state, valueType, source, offset, mask, l1_hint, l2_hint,
        l3_hint, layout, /*contiguity=*/nullptr);
}
 
//===----------------------------------------------------------------------===//
// XeGPU_StoreScatterOp
//===----------------------------------------------------------------------===//
LogicalResult StoreScatterOp::verify() {
  auto maskTy = getMaskType();
  auto valueTy = getValueType();
 
  if (!isWriteHintOrNone(getL1HintAttr()))
    return emitOpError("invalid l1_hint: ") << getL1HintAttr();
 
  if (!isWriteHintOrNone(getL2HintAttr()))
    return emitOpError("invalid l2_hint: ") << getL2HintAttr();
 
  if (!isWriteHintOrNone(getL3HintAttr()))
    return emitOpError("invalid l3_hint: ") << getL3HintAttr();
 
  auto destTy = getDestType();
  auto memTy = dyn_cast<MemRefType>(destTy);
 
  if (memTy && (getElementType() != memTy.getElementType()))
    return emitError() << "Value should have the same element type as MemRef.";
 
  if (auto layout = getAnchorLayout()) {
    if (!layout.isDistributable(getShapeOf(valueTy)))
      return emitOpError("Value shape is not distributable with the layout");
  }
 
  auto offsetsTy = getOffsets().getType();
  if (failed(isValidContiguity(getContiguity(), offsetsTy,
                               [&]() { return emitOpError(); })))
    return failure();
  return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy,
                                          [&]() { return emitOpError(); });
}
 
void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
                           Value value, Value dest,
                           ArrayRef<OpFoldResult> offsets, Value mask,
                           xegpu::CachePolicyAttr l1_hint,
                           xegpu::CachePolicyAttr l2_hint,
                           xegpu::CachePolicyAttr l3_hint) {
  auto loc = dest.getLoc();
  int64_t size = static_cast<int64_t>(offsets.size());
  auto type = VectorType::get(size, builder.getIndexType());
  auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
  auto offset = vector::FromElementsOp::create(builder, loc, type, values);
 
  // Call the correct builder overload that does not expect result types.
  build(builder, state, value, dest, offset, mask, l1_hint, l2_hint, l3_hint,
        /*anchor_layout=*/nullptr, /*contiguity=*/nullptr);
}
 
void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
                           Value value, Value dest,
                           ArrayRef<OpFoldResult> offsets, Value mask,
                           xegpu::CachePolicyAttr l1_hint,
                           xegpu::CachePolicyAttr l2_hint,
                           xegpu::CachePolicyAttr l3_hint,
                           DistributeLayoutAttr layout) {
  auto loc = dest.getLoc();
  int64_t size = static_cast<int64_t>(offsets.size());
  auto type = VectorType::get(size, builder.getIndexType());
  auto values = getValueOrCreateConstantIndexOp(builder, loc, offsets);
  auto offset = vector::FromElementsOp::create(builder, loc, type, values);
 
  // Call the correct builder overload that does not expect result types.
  build(builder, state, value, dest, offset, mask, l1_hint, l2_hint, l3_hint,
        layout, /*contiguity=*/nullptr);
}
 
//===----------------------------------------------------------------------===//
// DPAS Common Verification Helpers
//===----------------------------------------------------------------------===//
 
// Helper to verify layout distributability for a value
static LogicalResult
verifyLayoutDistributable(Operation *op,
                          std::optional<DistributeLayoutAttr> layout,
                          ArrayRef<int64_t> shape, StringRef operandName) {
  if (layout && !layout->isDistributable(
                    SmallVector<int64_t>(shape.begin(), shape.end())))
    return op->emitOpError(operandName)
           << " shape is not distributable with the layout";
  return success();
}
 
// Helper to verify M, N, K dimensions match between A, B, and result matrices
static LogicalResult verifyDpasDimensions(Operation *op,
                                          ArrayRef<int64_t> aShape,
                                          ArrayRef<int64_t> bShape,
                                          ArrayRef<int64_t> resShape) {
 
  auto aRank = aShape.size();
  auto bRank = bShape.size();
  auto resRank = resShape.size();
  if (aRank == 1 && bRank == 1 && resRank == 1)
    return success();
 
  // A must be at least 2D, B must be 2D or 3D (innermost dims), result at
  // least 2D.
  if (aRank < 2)
    return op->emitOpError("A operand must be at least a 2D vector.");
  if (bRank < 2)
    return op->emitOpError("B operand must be at least a 2D vector.");
  if (resRank < 2)
    return op->emitOpError("Result must be at least a 2D vector.");
 
  // FIXME: B may have one extra trailing dim for VNNI packing
  // (B[batch..., K/vnni, N, vnni]). We plan to drop VNNI packing support, so
  // rather than properly verifying the packed dimensions, we simply accept
  // the packed form here and skip the detailed verification. This branch
  // should be removed once VNNI packing support is dropped.
  if (bRank == aRank + 1)
    return success();
 
  // All operands have the same rank. They share the same batch dimensions,
  // with the last two dims being the core matmul dims: A[batch..., M, K],
  // B[batch..., K, N], result[batch..., M, N].
  if (aRank != bRank || aRank != resRank)
    return op->emitOpError("Rank mismatch among A, B, and result.");
 
  int64_t batchRank = aRank - 2;
 
  // Verify batch dimensions match.
  for (int64_t i = 0; i < batchRank; ++i) {
    if (aShape[i] != resShape[i])
      return op->emitOpError("Batch dimension mismatch at dim ")
             << i << ": A has " << aShape[i] << " but result has "
             << resShape[i] << ".";
    if (aShape[i] != bShape[i])
      return op->emitOpError("Batch dimension mismatch at dim ")
             << i << ": A has " << aShape[i] << " but B has " << bShape[i]
             << ".";
  }
 
  // Core matmul dimensions (last two dims of each operand).
  int64_t aM = aShape[batchRank];
  int64_t aK = aShape[batchRank + 1];
  int64_t bK = bShape[batchRank];
  int64_t bN = bShape[batchRank + 1];
  int64_t resM = resShape[batchRank];
  int64_t resN = resShape[batchRank + 1];
 
  // Verify K dimension match between A and B
  if (bK != aK)
    return op->emitOpError("K-dimension mismatch: A has K=")
           << aK << " but B has K=" << bK << ".";
 
  // Verify M dimension match between A and result
  if (aM != resM)
    return op->emitOpError("M-dimension mismatch: A has M=")
           << aM << " but result has M=" << resM << ".";
 
  // Verify N dimension match between B and result
  if (bN != resN)
    return op->emitOpError("N-dimension mismatch: B has N=")
           << bN << " but result has N=" << resN << ".";
 
  return success();
}
 
// Helper to verify accumulator matches result type
static LogicalResult verifyDpasAccumulator(Operation *op, Type accType,
                                           Type resultType) {
  if (accType != resultType)
    return op->emitOpError("Accumulator type must match result type.");
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_DpasOp
//===----------------------------------------------------------------------===//
LogicalResult DpasOp::verify() {
  auto lhsShape = getLhsType().getShape();
  auto rhsShape = getRhsType().getShape();
  auto resShape = getResultType().getShape();
 
  // Verify layout distributability
  if (failed(
          verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
    return failure();
  if (failed(verifyLayoutDistributable(*this, getLayoutA(), lhsShape, "A")))
    return failure();
  if (failed(verifyLayoutDistributable(*this, getLayoutB(), rhsShape, "B")))
    return failure();
 
  // Verify accumulator if present
  if (getAcc() &&
      failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
    return failure();
 
  return verifyDpasDimensions(*this, lhsShape, rhsShape, resShape);
}
 
//===----------------------------------------------------------------------===//
// XeGPU_ConvertLayoutOp
//===----------------------------------------------------------------------===//
LogicalResult ConvertLayoutOp::verify() {
  auto resLayout = getTargetLayout();
  if (!resLayout)
    return emitOpError("expected target layout.");
  auto srcLayout = getEffectiveInputLayout();
 
  // both input and target layouts should be WgLayout or SgLayout at the same
  // time.
  if ((!srcLayout.isForWorkgroup() || !resLayout.isForWorkgroup()) &&
      (!srcLayout.isForSubgroup() || !resLayout.isForSubgroup()))
    return emitOpError("expected input layout and target layout be WgLayout or "
                       "SgLayout at the same time.");
 
  Type srcType = getSource().getType();
  if (llvm::isa<VectorType>(srcType)) {
    SmallVector<int64_t> shape(llvm::cast<VectorType>(srcType).getShape());
    if (!srcLayout.isDistributable(shape))
      return emitOpError(
          "invalid input layout, data cannot be evenly distributed.");
 
    if (!resLayout.isDistributable(std::move(shape)))
      return emitOpError(
          "invalid target layout, data cannot be evenly distributed.");
  }
  return mlir::success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_LoadMatrixOp
//===----------------------------------------------------------------------===//
void LoadMatrixOp::build(OpBuilder &builder, OperationState &state, Type res,
                         TypedValue<MemDescType> memDesc,
                         llvm::ArrayRef<OpFoldResult> offsets,
                         DistributeLayoutAttr layout) {
  llvm::SmallVector<Value> dynamicOffsets;
  llvm::SmallVector<int64_t> staticOffsets;
  dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
  auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);
  // Call the generated builder with all parameters (including optional ones as
  // nullptr/empty)
  build(builder, state, res, memDesc, dynamicOffsets, staticOffsetsAttr,
        /*subgroup_block_io=*/nullptr, layout);
}
 
LogicalResult LoadMatrixOp::verify() {
 
  auto resTy = dyn_cast<VectorType>(getRes().getType());
  UnitAttr subgroup_block_io = getSubgroupBlockIoAttr();
  MemDescType mdescTy = getMemDesc().getType();
 
  return IsValidMatrixOpParams(resTy, mdescTy, subgroup_block_io,
                               getLayoutAttr(), [&]() { return emitError(); });
}
 
//===----------------------------------------------------------------------===//
// XeGPU_StoreMatrixOp
//===----------------------------------------------------------------------===//
void StoreMatrixOp::build(OpBuilder &builder, OperationState &state, Value data,
                          TypedValue<MemDescType> memDesc,
                          llvm::ArrayRef<OpFoldResult> offsets,
                          DistributeLayoutAttr layout) {
  llvm::SmallVector<Value> dynamicOffsets;
  llvm::SmallVector<int64_t> staticOffsets;
  dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
  auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);
  build(builder, state, data, memDesc, dynamicOffsets, staticOffsetsAttr,
        /*subgroup_block_io=*/nullptr, layout);
}
 
LogicalResult StoreMatrixOp::verify() {
 
  auto dataTy = dyn_cast<VectorType>(getData().getType());
  UnitAttr subgroup_block_io = getSubgroupBlockIoAttr();
  MemDescType mdescTy = getMemDesc().getType();
  return IsValidMatrixOpParams(dataTy, mdescTy, subgroup_block_io,
                               getLayoutAttr(), [&]() { return emitError(); });
}
 
//===----------------------------------------------------------------------===//
// XeGPU_TruncfOp
//===----------------------------------------------------------------------===//
 
LogicalResult TruncfOp::verify() {
  auto sourceVecType = dyn_cast<VectorType>(getSource().getType());
  auto resultVecType = dyn_cast<VectorType>(getResult().getType());
 
  if (sourceVecType.getElementTypeBitWidth() <=
      resultVecType.getElementTypeBitWidth())
    return emitOpError("input type must be wider than result type.");
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_LaneShuffleOp
//===----------------------------------------------------------------------===//
 
LogicalResult LaneShuffleOp::verify() {
  // With a single element per lane there is nothing to re-distribute, so the
  // operation would be a no-op.
  if (getSourceType().getNumElements() < 2)
    return emitOpError("requires a source vector with at least 2 elements.");
 
  return success();
}
 
OpFoldResult LaneShuffleOp::fold(FoldAdaptor adaptor) {
  // The two modes are exact inverses, so a pack feeding an unpack (or vice
  // versa) restores the original fragments.
  auto producer = getSource().getDefiningOp<LaneShuffleOp>();
  if (producer && producer.getMode() != getMode())
    return producer.getSource();
 
  return {};
}
 
//===----------------------------------------------------------------------===//
// XeGPU_DpasMxOp
//===----------------------------------------------------------------------===//
 
LogicalResult DpasMxOp::verify() {
  auto aShape = getAType().getShape();
  auto bShape = getBType().getShape();
  auto resShape = getResultType().getShape();
 
  // Verify layout distributability for A, B, and result
  if (failed(
          verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
    return failure();
  if (failed(verifyLayoutDistributable(*this, getLayoutA(), aShape, "A")))
    return failure();
  if (failed(verifyLayoutDistributable(*this, getLayoutB(), bShape, "B")))
    return failure();
 
  // Verify accumulator if present
  if (getAcc() &&
      failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
    return failure();
 
  // Verify M, N, K dimensions
  if (failed(verifyDpasDimensions(*this, aShape, bShape, resShape)))
    return failure();
 
  // Determine batch rank from A operand.
  int64_t aBatchRank = aShape.size() - 2;
 
  // Validate scale_a if present
  if (getScaleA()) {
    auto scaleAVecType = dyn_cast<VectorType>(getScaleAType());
    // Only validate if scale is a vector (scalars are always valid)
    if (scaleAVecType && scaleAVecType.getRank() > 1) {
      auto scaleAShape = scaleAVecType.getShape();
 
      if (scaleAVecType.getRank() < 2)
        return emitOpError("Scale A must be at least a 2D vector when not a "
                           "scalar.");
 
      // Verify layout distributability for scale_a
      if (failed(verifyLayoutDistributable(*this, getLayoutAScale(),
                                           scaleAShape, "ScaleA")))
        return failure();
 
      // Validate M dimension: scale_a's M must match A's M (last-1 dim)
      if (scaleAShape[scaleAShape.size() - 2] != aShape[aBatchRank])
        return emitOpError("Scale A M dimension [")
               << scaleAShape[scaleAShape.size() - 2]
               << "] must match A M dimension [" << aShape[aBatchRank] << "].";
    }
  }
 
  // Validate scale_b if present
  if (getScaleB()) {
    auto scaleBVecType = dyn_cast<VectorType>(getScaleBType());
    // Only validate if scale is a vector (scalars are always valid)
    if (scaleBVecType && scaleBVecType.getRank() > 1) {
      auto scaleBShape = scaleBVecType.getShape();
 
      if (scaleBVecType.getRank() < 2)
        return emitOpError("Scale B must be at least a 2D vector when not a "
                           "scalar.");
 
      // Verify layout distributability for scale_b
      if (failed(verifyLayoutDistributable(*this, getLayoutBScale(),
                                           scaleBShape, "ScaleB")))
        return failure();
 
      // Validate N dimension: scale_b's N (last dim) must match B's N (last
      // dim)
      if (scaleBShape.back() != bShape.back())
        return emitOpError("Scale B N dimension [")
               << scaleBShape.back() << "] must match B N dimension ["
               << bShape.back() << "].";
    }
  }
 
  // Validate scale K dimension compatibility if both scales are present and
  // vectors
  if (getScaleA() && getScaleB()) {
    auto scaleAVecType = dyn_cast<VectorType>(getScaleAType());
    auto scaleBVecType = dyn_cast<VectorType>(getScaleBType());
 
    if (scaleAVecType && scaleBVecType && scaleAVecType.getRank() > 1 &&
        scaleBVecType.getRank() > 1) {
      auto scaleAShape = scaleAVecType.getShape();
      auto scaleBShape = scaleBVecType.getShape();
 
      // Validate scale K dimension compatibility: scale_a's last dim must
      // match scale_b's second-to-last dim
      if (scaleAShape.back() != scaleBShape[scaleBShape.size() - 2])
        return emitOpError("Scale K dimension mismatch: scale_a has K=")
               << scaleAShape.back()
               << " but scale_b has K=" << scaleBShape[scaleBShape.size() - 2]
               << ".";
    }
  }
 
  return success();
}
 
namespace mlir {
#include <mlir/Dialect/XeGPU/IR/XeGPUAttrInterface.cpp.inc>
} // namespace mlir
#include <mlir/Dialect/XeGPU/IR/XeGPUEnums.cpp.inc>
#define GET_OP_CLASSES
#include <mlir/Dialect/XeGPU/IR/XeGPU.cpp.inc>