XeGPUDialect.cpp

Go to the documentation of this file.

//===- XeGPUDialect.cpp - MLIR XeGPU dialect 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/Affine/Utils.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/DialectImplementation.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
 
using std::optional;
 
namespace mlir {
namespace xegpu {
 
void XeGPUDialect::initialize() {
  addTypes<
#define GET_TYPEDEF_LIST
#include <mlir/Dialect/XeGPU/IR/XeGPUTypes.cpp.inc>
      >();
  addOperations<
#define GET_OP_LIST
#include <mlir/Dialect/XeGPU/IR/XeGPU.cpp.inc>
      >();
  addAttributes<
#define GET_ATTRDEF_LIST
#include <mlir/Dialect/XeGPU/IR/XeGPUAttrs.cpp.inc>
      >();
}
#define GET_OP_INTERFACE_CLASSES
#include "mlir/Dialect/XeGPU/IR/XeGPUOpInterface.cpp.inc"
 
// A `srcShape` consists of N distribution units, each being `subShapesLayout` x
// `subShape`. A `delinearizedId` is used to identify a particular `subShape`
// within each distribution unit.
// Example:
// WG data is 128x256. SG data is 16x32, in 4x2 layout, this gives a
// distribution unit of shape 64x64, we have 2x4 such distribution units.
// `delinearizedId` is used to identify a 16x32 of a subgroup in each
// distribution unit.
static SmallVector<SmallVector<Value>>
genCoordinates(OpBuilder &builder, Location loc,
               SmallVector<Value> delinearizedId,
               ArrayRef<int64_t> subShapesLayout, ArrayRef<int64_t> subShape,
               ArrayRef<int64_t> srcShape) {
  SmallVector<SmallVector<Value>> coordinates;
 
  // A distribution unit must be less than or equal to `srcShape`
  SmallVector<int64_t> distUnitShape = llvm::map_to_vector(
      llvm::zip_equal(srcShape,
                      computeElementwiseMul(subShapesLayout, subShape)),
      [](const auto &t) { return std::min(std::get<0>(t), std::get<1>(t)); });
 
  // Get the offset of `subShape` within a distribution unit.
  SmallVector<Value> distUnitLocalOffset = llvm::map_to_vector(
      llvm::zip(delinearizedId, subShape), [&](const auto &t) -> Value {
        return builder.createOrFold<arith::MulIOp>(
            loc, std::get<0>(t),
            builder.createOrFold<arith::ConstantIndexOp>(loc, std::get<1>(t)));
      });
 
  // For each dist unit
  for (SmallVector<int64_t> unitOffs :
       StaticTileOffsetRange(srcShape, distUnitShape)) {
    // Get dist unit offset within `srcShape`.
    SmallVector<Value> base =
        llvm::map_to_vector(unitOffs, [&](int64_t d) -> Value {
          return arith::ConstantIndexOp::create(builder, loc, d);
        });
    // Calculate `subShape` offset within `srcShape`.
    SmallVector<Value> adds =
        llvm::map_to_vector(llvm::zip_equal(base, distUnitLocalOffset),
                            [&](const auto &t) -> Value {
                              return builder.createOrFold<arith::AddIOp>(
                                  loc, std::get<0>(t), std::get<1>(t));
                            });
    // Do not go beyond `srcShape` bounds.
    SmallVector<Value> mods = llvm::map_to_vector(
        llvm::zip_equal(adds, srcShape), [&](const auto &t) -> Value {
          return builder.createOrFold<arith::RemUIOp>(
              loc, std::get<0>(t),
              arith::ConstantIndexOp::create(builder, loc, std::get<1>(t)));
        });
 
    coordinates.push_back(mods);
  }
  return coordinates;
}
 
static SmallVector<SmallVector<int64_t>> genStaticCoordinates(
    llvm::ArrayRef<int64_t> canonicalIds, llvm::ArrayRef<int64_t> layout,
    llvm::ArrayRef<int64_t> subShape, llvm::ArrayRef<int64_t> shape) {
  // Compute distribution unit shape (clamped to srcShape).
  SmallVector<int64_t> distUnitShape(shape.size());
  for (size_t i = 0; i < shape.size(); ++i)
    distUnitShape[i] = std::min(shape[i], layout[i] * subShape[i]);
 
  // Compute local offset of this ID within a distribution unit.
  SmallVector<int64_t> localOffset(shape.size());
  for (size_t i = 0; i < shape.size(); ++i)
    localOffset[i] = canonicalIds[i] * subShape[i];
 
  // Enumerate all distribution units and compute coordinates.
  SmallVector<SmallVector<int64_t>> coordinates;
  for (SmallVector<int64_t> unitOffs :
       StaticTileOffsetRange(shape, distUnitShape)) {
    SmallVector<int64_t> coord(shape.size());
    for (size_t i = 0; i < shape.size(); ++i)
      coord[i] = (unitOffs[i] + localOffset[i]) % shape[i];
    coordinates.push_back(coord);
  }
  return coordinates;
}
 
/// Expands per-distribution-unit block-start coordinates into the full list of
/// element coordinates each block covers: every element of the `subShape`-sized
/// region (row-major) offset by the block start. Comparing these instead of the
/// bare block starts lets layouts that differ only in `lane_data` blocking, but
/// own the same elements in the same order, be recognized as equivalent.
static SmallVector<SmallVector<int64_t>>
expandBlockCoords(ArrayRef<SmallVector<int64_t>> blockStarts,
                  ArrayRef<int64_t> subShape) {
  SmallVector<int64_t> unitTile(subShape.size(), 1);
  SmallVector<SmallVector<int64_t>> expanded;
  for (const SmallVector<int64_t> &start : blockStarts) {
    for (SmallVector<int64_t> off : StaticTileOffsetRange(subShape, unitTile)) {
      SmallVector<int64_t> coord(start.size());
      for (size_t i = 0; i < start.size(); ++i)
        coord[i] = start[i] + off[i];
      expanded.push_back(std::move(coord));
    }
  }
  return expanded;
}
 
/// Returns true if `self` and `other` distribute `shape` identically at
/// `level`: every id in `[0, size)` owns the same coordinates under both.
///
/// At the Lane level, layouts that pack `lane_data` differently can still own
/// the same per-lane elements in the same order; their block starts differ but
/// the expanded per-element coordinates match. So block starts are expanded
/// (via `expandBlockCoords`) before comparing, but only when it can change the
/// result (Lane level with differing `lane_data`) - otherwise comparing the
/// cheaper block starts is already exact.
///
/// TODO: Extend the same handling to the Subgroup level (sg_data repacks).
static bool compareDistributedCoords(xegpu::DistributeLayoutAttr self,
                                     const xegpu::DistributeLayoutAttr &other,
                                     ArrayRef<int64_t> shape,
                                     xegpu::LayoutKind level, int64_t size) {
  bool expandCoords =
      level == xegpu::LayoutKind::Lane &&
      self.getEffectiveLaneDataAsInt() != other.getEffectiveLaneDataAsInt();
  SmallVector<int64_t> selfSubShape, otherSubShape;
  if (expandCoords) {
    selfSubShape = self.getEffectiveLaneDataAsInt();
    otherSubShape = other.getEffectiveLaneDataAsInt();
  }
  for (int64_t id : llvm::seq<int64_t>(0, size)) {
    auto coords = self.computeStaticDistributedCoords(id, shape);
    auto otherCoords = other.computeStaticDistributedCoords(id, shape);
    if (expandCoords) {
      coords = expandBlockCoords(coords, selfSubShape);
      otherCoords = expandBlockCoords(otherCoords, otherSubShape);
    }
    if (coords != otherCoords)
      return false;
  }
  return true;
}
 
// Checks if the given memref type represents shared local memory (SLM).
bool XeGPUDialect::isSharedMemory(const MemRefType &memrefTy) {
  Attribute attr = memrefTy.getMemorySpace();
  if (!attr)
    return false; // Default memory space is not shared local memory
  if (auto intAttr = llvm::dyn_cast_if_present<IntegerAttr>(attr))
    return intAttr.getInt() == 3;
  if (auto memrefSpace = llvm::dyn_cast_if_present<MemorySpaceAttr>(attr))
    return memrefSpace.getValue() == MemorySpace::SLM;
  if (auto xevmSpace = llvm::dyn_cast_if_present<xevm::AddrSpaceAttr>(attr))
    return xevmSpace.getValue() == xevm::AddrSpace::SHARED;
  return gpu::GPUDialect::isWorkgroupMemoryAddressSpace(attr);
}
 
//===----------------------------------------------------------------------===//
// XeGPU_BlockTensorDescAttr
//===----------------------------------------------------------------------===//
BlockTensorDescAttr BlockTensorDescAttr::get(mlir::MLIRContext *context,
                                             xegpu::MemorySpace memory_space,
                                             int array_length,
                                             bool boundary_check) {
  auto scopeAttr = MemorySpaceAttr::get(context, memory_space);
  auto lengthAttr =
      IntegerAttr::get(IntegerType::get(context, 64), array_length);
  auto boundaryAttr = BoolAttr::get(context, boundary_check);
  return Base::get(context, scopeAttr, lengthAttr, boundaryAttr);
}
 
bool BlockTensorDescAttr::hasDefaultsOnly() {
  return getMemorySpace().getValue() == xegpu::MemorySpace::Global &&
         getArrayLength().getInt() == 1 && getBoundaryCheck().getValue();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_LayoutAttr
//===----------------------------------------------------------------------===//
LogicalResult
LayoutAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
                   DenseI32ArrayAttr sg_layout, DenseI32ArrayAttr sg_data,
                   DenseI32ArrayAttr inst_data, DenseI32ArrayAttr lane_layout,
                   DenseI32ArrayAttr lane_data, DenseI32ArrayAttr order) {
 
  // Special case for store_matrix
  if (!sg_layout && !inst_data && !lane_layout)
    return success();
 
  // generate code to check sg_laout, inst_data and lane_layout having the same
  // rank if they are not null.
 
  if (sg_layout && inst_data && sg_layout.size() != inst_data.size()) {
    return emitError()
           << "expected sg_layout and inst_data to have the same rank";
  }
 
  if (sg_layout && lane_layout && sg_layout.size() != lane_layout.size()) {
    return emitError()
           << "expected sg_layout and lane_layout to have the same rank";
  }
 
  if (inst_data && lane_layout && inst_data.size() != lane_layout.size()) {
    return emitError() << "expected inst_data and lane_layout to have the same "
                          "rank, got inst_data "
                       << inst_data.size() << ", lane_layout "
                       << lane_layout.size();
  }
 
  if ((sg_layout && !sg_data) || (!sg_layout && sg_data))
    return emitError() << "sg_layout and sg_data must be used together";
  if (sg_layout && sg_data && sg_layout.size() != sg_data.size())
    return emitError()
           << "expected sg_data and sg_layout to have the same rank";
 
  if ((lane_layout && !lane_data) || (!lane_layout && lane_data))
    return emitError() << "lane_layout and lane_data must be used together";
  if (lane_layout && lane_data && lane_layout.size() != lane_data.size())
    return emitError()
           << "expected lane_data and lane_layout to have the same rank";
 
  if (order) {
    if (!sg_layout && !lane_layout)
      return emitError()
             << "expected sg_layout/lane_layout being used with order";
 
    if (sg_layout && order.size() != sg_layout.size())
      return emitError()
             << "expected order and sg_layout to have the same rank";
 
    if (lane_layout && order.size() != lane_layout.size())
      return emitError()
             << "expected order and lane_layout to have the same rank";
  }
 
  return success();
}
 
FailureOr<SmallVector<Value>>
LayoutAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId) {
 
  SmallVector<int64_t> sgLayoutInt;
  if (isForWorkgroup()) {
    sgLayoutInt = getEffectiveSgLayoutAsInt();
  } else if (isForSubgroup()) {
    sgLayoutInt = getEffectiveLaneLayoutAsInt();
  } else {
    return failure();
  }
 
  DenseI32ArrayAttr orderAttr = getOrder();
 
  // Handle order attribute
  SmallVector<int64_t> order;
  if (orderAttr && !orderAttr.empty()) {
    order = llvm::map_to_vector(orderAttr.asArrayRef(), [](int32_t idx) {
      return static_cast<int64_t>(idx);
    });
  } else {
    // Default order: [1, 0] for 2D (row-major), [2, 1, 0] for 3D, etc.
    order = llvm::to_vector(
        llvm::reverse(llvm::seq<int64_t>(0, sgLayoutInt.size())));
  }
 
  if (order.size() != sgLayoutInt.size()) {
    return failure();
  }
 
  SmallVector<Value> result(sgLayoutInt.size());
  Value remaining = linearId;
 
  /// Process dimensions in the order they appear in the order array
  /// The first dimension in order is the fastest-changing
  ///
  /// Example walkthrough for linearId=22, sgLayout=[2,4,4], order=[2,1,0]:
  ///
  /// Initial: remaining=22, dimIdx = order[i], dimSize = sgLayout[dimIdx],
  /// result=[?,?,?]
  ///
  /// i=0 (process columns, dimIdx=2, dimSize=4):
  ///   result[2] = 22 % 4 = 2  (column coordinate)
  ///   remaining = 22 / 4 = 5  (5 complete groups of 4 columns processed)
  ///
  /// i=1 (process rows, dimIdx=1, dimSize=4):
  ///   result[1] = 5 % 4 = 1   (row coordinate)
  ///   remaining = 5 / 4 = 1   (1 complete group of 4 rows processed)
  ///
  /// i=2 (process layers, dimIdx=0, dimSize=2):
  ///   result[0] = 1 % 2 = 1   (layer coordinate)
  ///   (no remaining update - last iteration)
  ///
  /// Final result: [1,1,2] = Layer 1, Row 1, Column 2
  for (size_t i = 0; i < order.size(); ++i) {
    int64_t dimIdx = order[i];
    int64_t dimSize = sgLayoutInt[dimIdx];
 
    Value dimSizeVal =
        builder.createOrFold<arith::ConstantIndexOp>(loc, dimSize);
 
    /// Extract the coordinate for this dimension using modulo operation
    /// This gives us "how far within this dimension" we are
    /// e.g., linearId=22, dimSize=4: 22 % 4 = 2 (we're at position 2 within
    /// this dimension)
    result[dimIdx] =
        builder.createOrFold<arith::RemUIOp>(loc, remaining, dimSizeVal);
 
    /// Update remaining for the next dimension by removing what we've already
    /// processed. Division tells us "how many complete groups of this dimension
    /// we've gone through" e.g., linearId=22, dimSize=4: 22 / 4 = 5 (we've
    /// completed 5 groups of 4) Skip this for the last iteration since there's
    /// no next dimension to process
    if (i < order.size() - 1) {
      remaining =
          builder.createOrFold<arith::DivUIOp>(loc, remaining, dimSizeVal);
    }
  }
  return result;
}
 
/// Implements DistributeLayoutAttr::computeDistributedCoords to generate
/// instructions for computing multi-dimensional offsets when distributed by
/// LayoutAttr.
FailureOr<SmallVector<SmallVector<Value>>>
LayoutAttr::computeDistributedCoords(OpBuilder &builder, Location loc,
                                     Value linearId, ArrayRef<int64_t> shape) {
  SmallVector<int64_t> layout;
  SmallVector<int64_t> subShape;
  if (isForWorkgroup()) {
    layout = getEffectiveSgLayoutAsInt();
    subShape = getEffectiveSgDataAsInt();
  } else if (isForSubgroup()) {
    layout = getEffectiveLaneLayoutAsInt();
    subShape = getEffectiveLaneDataAsInt();
  } else {
    return failure();
  }
  assert(!subShape.empty() && "sgdata or lanedata cannot be empty for "
                              "distributed coordinates computation");
 
  // delinearize Ids
  auto maybeIds = delinearizeId(builder, loc, linearId);
  if (failed(maybeIds))
    return failure();
  SmallVector<Value> ids = *maybeIds;
 
  return genCoordinates(builder, loc, ids, layout, subShape, shape);
}
 
bool LayoutAttr::isEqualTo(const xegpu::DistributeLayoutAttr &other) {
  if (dyn_cast<xegpu::SliceAttr>(other))
    return false;
 
  return *this == dyn_cast<xegpu::LayoutAttr>(other);
}
 
/// Implements DistributeLayoutAttr::computeStaticDistributedCoords to
/// compute multi-dimensional offsets for a given linear ID when distributed by
/// LayoutAttr.
SmallVector<SmallVector<int64_t>>
LayoutAttr::computeStaticDistributedCoords(int64_t linearId,
                                           ArrayRef<int64_t> shape) {
  SmallVector<int64_t> layoutVec;
  SmallVector<int64_t> subShape;
  SmallVector<int64_t> instData;
  if (isForWorkgroup()) {
    layoutVec = getEffectiveSgLayoutAsInt();
    subShape = getEffectiveSgDataAsInt();
  } else if (isForSubgroup()) {
    instData = getEffectiveInstDataAsInt();
    layoutVec = getEffectiveLaneLayoutAsInt();
    subShape = getEffectiveLaneDataAsInt();
  }
  if (!instData.empty()) {
    linearId = 0;
    subShape = instData;
  }
  assert(!subShape.empty() && "sgdata or lanedata cannot be empty");
 
  // Delinearize the linear ID using the order attribute.
  SmallVector<int64_t> order = getEffectiveOrderAsInt();
  SmallVector<int64_t> delinearizedId(layoutVec.size());
  int64_t remaining = linearId;
  for (size_t i = 0; i < order.size(); ++i) {
    int64_t dimIdx = order[i];
    delinearizedId[dimIdx] = remaining % layoutVec[dimIdx];
    remaining = remaining / layoutVec[dimIdx];
  }
 
  return genStaticCoordinates(delinearizedId, layoutVec, subShape, shape);
}
 
// set the layout for unit dims: sg_data, inst_data and lane_data to 1
DistributeLayoutAttr
LayoutAttr::setUnitDimData(SmallVector<int64_t> unitDims) const {
  auto sgDataOpt = getSgData();
  auto instDataOpt = getInstData();
  auto laneDataOpt = getLaneData();
 
  SmallVector<int32_t> sgData;
  SmallVector<int32_t> instData;
  SmallVector<int32_t> laneData;
 
  if (sgDataOpt)
    sgData = llvm::to_vector(sgDataOpt.asArrayRef());
 
  if (instDataOpt)
    instData = llvm::to_vector(instDataOpt.asArrayRef());
 
  if (laneDataOpt)
    laneData = llvm::to_vector(laneDataOpt.asArrayRef());
 
  for (auto dim : unitDims) {
    if (dim < static_cast<int64_t>(sgData.size()))
      sgData[dim] = 1;
    if (dim < static_cast<int64_t>(instData.size()))
      instData[dim] = 1;
    if (dim < static_cast<int64_t>(laneData.size()))
      laneData[dim] = 1;
  }
 
  return LayoutAttr::get(
      getContext(), getSgLayout(),
      sgData.empty() ? DenseI32ArrayAttr()
                     : DenseI32ArrayAttr::get(getContext(), sgData),
      instData.empty() ? DenseI32ArrayAttr()
                       : DenseI32ArrayAttr::get(getContext(), instData),
      getLaneLayout(),
      laneData.empty() ? DenseI32ArrayAttr()
                       : DenseI32ArrayAttr::get(getContext(), laneData),
      getOrder());
}
 
// set the layout for the sepcified unit dims: sg_lane and lane_layout to 1
DistributeLayoutAttr
LayoutAttr::setUnitDimLayout(SmallVector<int64_t> unitDims) const {
  auto sgLayoutOpt = getSgLayout();
  auto laneLayoutOpt = getLaneLayout();
 
  SmallVector<int32_t> sgLayout;
  SmallVector<int32_t> laneLayout;
 
  if (sgLayoutOpt)
    sgLayout = llvm::to_vector(sgLayoutOpt.asArrayRef());
  if (laneLayoutOpt)
    laneLayout = llvm::to_vector(laneLayoutOpt.asArrayRef());
 
  for (auto dim : unitDims) {
    if (dim < static_cast<int64_t>(sgLayout.size()))
      sgLayout[dim] = 1;
    if (dim < static_cast<int64_t>(laneLayout.size()))
      laneLayout[dim] = 1;
  }
 
  return LayoutAttr::get(
      getContext(),
      sgLayout.empty() ? DenseI32ArrayAttr()
                       : DenseI32ArrayAttr::get(getContext(), sgLayout),
      getSgData(), getInstData(),
      laneLayout.empty() ? DenseI32ArrayAttr()
                         : DenseI32ArrayAttr::get(getContext(), laneLayout),
      getLaneData(), getOrder());
}
 
// Derive a new layout with sg_data, inst_data and lane_data set to the
// specified values for the given dimension
DistributeLayoutAttr LayoutAttr::setDimData(int64_t dim, int64_t sgData,
                                            int64_t instData,
                                            int64_t laneData) {
 
  SmallVector<int64_t> sgDataVec = getEffectiveSgDataAsInt();
  SmallVector<int64_t> instDataVec = getEffectiveInstDataAsInt();
  SmallVector<int64_t> laneDataVec = getEffectiveLaneDataAsInt();
 
  if (dim < static_cast<int64_t>(sgDataVec.size()) && sgData != -1)
    sgDataVec[dim] = sgData;
  if (dim < static_cast<int64_t>(instDataVec.size()) && instData != -1)
    instDataVec[dim] = instData;
  if (dim < static_cast<int64_t>(laneDataVec.size()) && laneData != -1)
    laneDataVec[dim] = laneData;
 
  SmallVector<int32_t> sgDataVec32(sgDataVec.begin(), sgDataVec.end());
  SmallVector<int32_t> instDataVec32(instDataVec.begin(), instDataVec.end());
  SmallVector<int32_t> laneDataVec32(laneDataVec.begin(), laneDataVec.end());
 
  return LayoutAttr::get(
      getContext(), getSgLayout(),
      sgDataVec.empty() ? DenseI32ArrayAttr()
                        : DenseI32ArrayAttr::get(getContext(), sgDataVec32),
      instDataVec.empty() ? DenseI32ArrayAttr()
                          : DenseI32ArrayAttr::get(getContext(), instDataVec32),
      getLaneLayout(),
      laneDataVec.empty() ? DenseI32ArrayAttr()
                          : DenseI32ArrayAttr::get(getContext(), laneDataVec32),
      getOrder());
}
 
// Derive a new layout by removing dimensions.
// `dimGroup` specifies a group of dimensions to be removed in the derived
// layout.
DistributeLayoutAttr LayoutAttr::dropDims(SmallVector<int64_t> dimGroup) {
 
  SmallVector<int64_t> sgLayout = getEffectiveSgLayoutAsInt();
  SmallVector<int64_t> sgData = getEffectiveSgDataAsInt();
  SmallVector<int64_t> instData = getEffectiveInstDataAsInt();
  SmallVector<int64_t> laneLayout = getEffectiveLaneLayoutAsInt();
  SmallVector<int64_t> laneData = getEffectiveLaneDataAsInt();
  DenseI32ArrayAttr origOrderAttr = getOrder();
 
  SmallVector<int64_t> sortedDimGroup = dimGroup;
  llvm::sort(sortedDimGroup);
 
  for (auto dimIdx : llvm::reverse(sortedDimGroup)) {
    if (!sgLayout.empty()) {
      sgLayout.erase(sgLayout.begin() + dimIdx);
      sgData.erase(sgData.begin() + dimIdx);
    }
    if (!instData.empty())
      instData.erase(instData.begin() + dimIdx);
    if (!laneLayout.empty()) {
      laneLayout.erase(laneLayout.begin() + dimIdx);
      laneData.erase(laneData.begin() + dimIdx);
    }
  }
 
  // Only emit a new order attribute when the input had one, so that "no
  // order" inputs do not gain a synthetic default that would later trip
  // adjacency checks in collapseDims.
  SmallVector<int64_t> newOrder;
  if (origOrderAttr && !origOrderAttr.empty()) {
    SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
    for (int64_t d : origOrder) {
      if (llvm::is_contained(dimGroup, d))
        continue;
      int64_t offset =
          llvm::count_if(dimGroup, [&](int64_t s) { return s < d; });
      newOrder.push_back(d - offset);
    }
    if ((sgLayout.empty() && laneLayout.empty()) || newOrder.size() == 1)
      newOrder.clear();
  }
 
  auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
    if (v.empty())
      return nullptr;
    SmallVector<int32_t> v32(v.begin(), v.end());
    return DenseI32ArrayAttr::get(getContext(), v32);
  };
  auto droppedLayout = xegpu::LayoutAttr::get(
      getContext(), toAttr(sgLayout), toAttr(sgData), toAttr(instData),
      toAttr(laneLayout), toAttr(laneData), toAttr(newOrder));
  return droppedLayout;
}
 
// Derive a new layout by collapsing dimensions.
// `dimGroup` specifies a group of adjacent dimensions
// that are collapsed into a single dimension in the derived layout.
DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
 
  SmallVector<int64_t> sgLayout = getEffectiveSgLayoutAsInt();
  SmallVector<int64_t> sgData = getEffectiveSgDataAsInt();
  SmallVector<int64_t> instData = getEffectiveInstDataAsInt();
  SmallVector<int64_t> laneLayout = getEffectiveLaneLayoutAsInt();
  SmallVector<int64_t> laneData = getEffectiveLaneDataAsInt();
  SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
 
  SmallVector<int64_t> sortedDimGroup = dimGroup;
  llvm::sort(sortedDimGroup);
  // Dims must be adjacent in the walk order
  bool hasExplicitWalkOrder = getOrder() && !getOrder().empty();
  for (size_t dimIdx = 1; dimIdx < sortedDimGroup.size(); ++dimIdx) {
    int64_t prev = sortedDimGroup[dimIdx - 1];
    int64_t curr = sortedDimGroup[dimIdx];
 
    if (hasExplicitWalkOrder) {
      // Walk order matters if dims are not trivially collapsible
      if ((sgLayout.empty() || (sgLayout[prev] == 1 && sgLayout[curr] == 1)) &&
          (laneLayout.empty() ||
           (laneLayout[prev] == 1 && laneLayout[curr] == 1)))
        continue;
      if (std::abs(origOrder[prev] - origOrder[curr]) != 1)
        llvm::report_fatal_error(
            "dimensions being collapsed must be adjacent in order");
    } else if (curr - prev != 1)
      llvm::report_fatal_error("dimensions being collapsed must be adjacent");
  }
 
  int firstDim = sortedDimGroup.front();
 
  // collapse the dimensions in dimGroup into one dimension by multiplying their
  // sizes together
 
  if (!sgLayout.empty()) {
    int64_t collapsedSglayout = 1, collapsedSgData = 1;
    for (auto dimIdx : dimGroup) {
      collapsedSglayout *= sgLayout[dimIdx];
      collapsedSgData *= sgData[dimIdx];
    }
    for (auto dimIdx : llvm::reverse(sortedDimGroup)) {
      sgLayout.erase(sgLayout.begin() + dimIdx, sgLayout.begin() + dimIdx + 1);
      sgData.erase(sgData.begin() + dimIdx, sgData.begin() + dimIdx + 1);
    }
    sgLayout.insert(sgLayout.begin() + firstDim, collapsedSglayout);
    sgData.insert(sgData.begin() + firstDim, collapsedSgData);
  }
 
  if (!instData.empty()) {
    int64_t collapsedInstData = 1;
    for (auto dimIdx : dimGroup)
      collapsedInstData *= instData[dimIdx];
    for (auto dimIdx : llvm::reverse(sortedDimGroup))
      instData.erase(instData.begin() + dimIdx, instData.begin() + dimIdx + 1);
    instData.insert(instData.begin() + firstDim, collapsedInstData);
  }
 
  if (!laneLayout.empty()) {
    int64_t collapsedLaneLayout = 1, collapsedLaneData = 1;
    for (auto dimIdx : dimGroup) {
      collapsedLaneLayout *= laneLayout[dimIdx];
      collapsedLaneData *= laneData[dimIdx];
    }
    for (auto dimIdx : llvm::reverse(sortedDimGroup)) {
      laneLayout.erase(laneLayout.begin() + dimIdx,
                       laneLayout.begin() + dimIdx + 1);
      laneData.erase(laneData.begin() + dimIdx, laneData.begin() + dimIdx + 1);
    }
    laneLayout.insert(laneLayout.begin() + firstDim, collapsedLaneLayout);
    laneData.insert(laneData.begin() + firstDim, collapsedLaneData);
  }
 
  SmallVector<int64_t> newOrder;
  DenseI32ArrayAttr orderAttr = getOrder();
  if (orderAttr && !orderAttr.empty()) {
 
    for (auto dimIdx : llvm::reverse(sortedDimGroup)) {
      if (dimIdx != firstDim)
        origOrder.erase(origOrder.begin() + dimIdx);
    }
    // say we have orderVec = {5, 3, 2, 1, 0}
    // Create indices [0, 1, 2, 3, 4]
    SmallVector<size_t> indices =
        llvm::to_vector(llvm::seq<size_t>(0, origOrder.size()));
 
    // Sort indices based on corresponding values
    llvm::sort(indices,
               [&](size_t a, size_t b) { return origOrder[a] < origOrder[b]; });
 
    newOrder = llvm::to_vector(llvm::map_range(
        indices, [&](size_t i) { return static_cast<int64_t>(i); }));
  }
 
  auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
    if (v.empty())
      return nullptr;
    SmallVector<int32_t> v32(v.begin(), v.end());
    return DenseI32ArrayAttr::get(getContext(), v32);
  };
  auto collapsedLayout = xegpu::LayoutAttr::get(
      getContext(), toAttr(sgLayout), toAttr(sgData), toAttr(instData),
      toAttr(laneLayout), toAttr(laneData), toAttr(newOrder));
  return collapsedLayout;
}
 
// Derive a new layout by expanding a single dimension `dim` into multiple
// adjacent dimensions whose extents are given by `targetShape`.
//
// The expanded dims are row-major (innermost is fastest-varying), so every
// field is spread innermost-first. For sg and lane levels the data component
// (sg_data / lane_data) is distributed first and the layout component
// (sg_layout / lane_layout) then strides over the per-dim leftover; this
// shared inner-first sweep is what makes "wrap-around" distributions correct.
// When the data component spans the full extent (replicated/broadcast across
// subgroups or lanes), the layout component is capped by the full extent
// instead of the leftover. order entries past `dim` shift up by
// `targetShape.size() - 1`.
//
// Examples (only the affected dim shown; assume rank-1 input):
//   sg_layout=[8], sg_data=[512],   expandDim(0, [8,16,32])
//     -> sg_layout=[8,1,1],   sg_data=[1,16,32]
//   sg_layout=[4], sg_data=[4],     expandDim(0, [2,16])
//     -> sg_layout=[1,4],     sg_data=[1,4]      (wrap-around)
//   inst_data=[32], lane_layout=[16], lane_data=[1], expandDim(0, [8,16,32])
//     -> inst_data=[1,1,32], lane_layout=[1,1,16], lane_data=[1,1,1]
DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
                                           ArrayRef<int64_t> targetShape) {
  SmallVector<int64_t> sgLayout = getEffectiveSgLayoutAsInt();
  SmallVector<int64_t> sgData = getEffectiveSgDataAsInt();
  SmallVector<int64_t> instData = getEffectiveInstDataAsInt();
  SmallVector<int64_t> laneLayout = getEffectiveLaneLayoutAsInt();
  SmallVector<int64_t> laneData = getEffectiveLaneDataAsInt();
 
  int64_t origRank = getRank();
  int64_t expCount = static_cast<int64_t>(targetShape.size());
  assert(dim >= 0 && dim < origRank && "dim out of range");
  assert(expCount >= 1 && "targetShape must have at least one dim");
  int64_t newRank = origRank + expCount - 1;
 
  // Snapshot the per-dim values we need before any field is mutated by
  // splice() below; without this, computations that read e.g. `laneLayout[dim]`
  // after laneLayout has already been expanded would see the wrong value.
  int64_t origSgLayoutDim = sgLayout.empty() ? 1 : sgLayout[dim];
  int64_t origSgDataDim = sgData.empty() ? 1 : sgData[dim];
  int64_t origLaneLayoutDim = laneLayout.empty() ? 1 : laneLayout[dim];
  int64_t origLaneDataDim = laneData.empty() ? 1 : laneData[dim];
  int64_t origInstDataDim = instData.empty() ? 1 : instData[dim];
 
  // Spread `total` across the new dims (length expCount), innermost-first,
  // capped per dim by `dimSizeCap[i]`.
  auto spread = [&](int64_t total,
                    ArrayRef<int64_t> dimSizeCap) -> SmallVector<int64_t> {
    SmallVector<int64_t> out(expCount, 1);
    int64_t remaining = total;
    for (int64_t i = expCount - 1; i >= 0; --i) {
      if (remaining == 1)
        break;
      int64_t take = std::min(remaining, dimSizeCap[i]);
      assert(take > 0 && "expandDim distribution must not be zero");
      assert(remaining % take == 0 &&
             "expandDims must divide evenly across dims");
      out[i] = take;
      remaining /= take;
    }
    assert(remaining == 1 && "expandDims total must fit within target shape");
    return out;
  };
 
  // Splice `expanded` (length expCount) into `vec` at position `dim`,
  // replacing the single entry at `dim`.
  auto splice = [&](SmallVector<int64_t> &vec, ArrayRef<int64_t> expanded) {
    if (vec.empty())
      return;
    vec.erase(vec.begin() + dim);
    vec.insert(vec.begin() + dim, expanded.begin(), expanded.end());
  };
 
  bool hasSgLayout = !sgLayout.empty();
  bool hasSgData = !sgData.empty();
  bool hasLaneLayout = !laneLayout.empty();
  bool hasLaneData = !laneData.empty();
  bool hasInstData = !instData.empty();
 
  // sg_data first, then sg_layout into the per-dim leftover. When sg_data is
  // replicated across subgroups it spans the full extent, so sg_layout is
  // capped by targetShape itself instead of the leftover.
  bool sgDataReplicated =
      hasSgData && origSgDataDim == computeProduct(targetShape);
  SmallVector<int64_t> expSgData(expCount, 1);
  if (hasSgData) {
    expSgData = spread(origSgDataDim, targetShape);
    splice(sgData, expSgData);
  }
  SmallVector<int64_t> expSgLayout(expCount, 1);
  if (hasSgLayout) {
    SmallVector<int64_t> dimSizeCap(targetShape.begin(), targetShape.end());
    if (hasSgData && !sgDataReplicated)
      for (int64_t i = 0; i < expCount; ++i)
        dimSizeCap[i] /= expSgData[i];
    expSgLayout = spread(origSgLayoutDim, dimSizeCap);
    splice(sgLayout, expSgLayout);
  }
 
  // Per-sg extent feeding lane_layout / lane_data / inst_data.
  SmallVector<int64_t> perSgShape(targetShape.begin(), targetShape.end());
  if (hasSgLayout && !sgDataReplicated)
    for (int64_t i = 0; i < expCount; ++i)
      perSgShape[i] /= expSgLayout[i];
 
  // lane_data first, then lane_layout into the per-dim leftover. When lane_data
  // is replicated across lanes it spans the full per-sg extent, so lane_layout
  // is capped by perSgShape itself instead of the leftover.
  bool laneDataReplicated =
      hasLaneData && origLaneDataDim == computeProduct(perSgShape);
  SmallVector<int64_t> expLaneLayout(expCount, 1);
  SmallVector<int64_t> expLaneData(expCount, 1);
  if (hasLaneData)
    expLaneData = spread(origLaneDataDim, perSgShape);
  if (hasLaneLayout) {
    SmallVector<int64_t> dimSizeCap(perSgShape.begin(), perSgShape.end());
    if (hasLaneData && !laneDataReplicated)
      for (int64_t i = 0; i < expCount; ++i)
        dimSizeCap[i] /= expLaneData[i];
    expLaneLayout = spread(origLaneLayoutDim, dimSizeCap);
  }
  if (hasLaneData)
    splice(laneData, expLaneData);
  if (hasLaneLayout)
    splice(laneLayout, expLaneLayout);
 
  // inst_data: seed each dim from the per-lane atom
  // `laneLayout[i]*laneData[i]`, spread the remaining factor over the leftover,
  // then scale back by the atom. Without lane info, spread inst_data directly
  // over the per-sg extent.
  if (hasInstData) {
    SmallVector<int64_t> expInstData;
    if (!hasLaneLayout || !hasLaneData) {
      expInstData = spread(origInstDataDim, perSgShape);
    } else {
      int64_t laneAtom = origLaneLayoutDim * origLaneDataDim;
      SmallVector<int64_t> atom(expCount, 1);
      SmallVector<int64_t> dimSizeCap(expCount, 1);
      for (int64_t i = 0; i < expCount; ++i) {
        atom[i] = expLaneLayout[i] * expLaneData[i];
        dimSizeCap[i] = perSgShape[i] / atom[i];
      }
      expInstData = spread(origInstDataDim / laneAtom, dimSizeCap);
      for (int64_t i = 0; i < expCount; ++i)
        expInstData[i] *= atom[i];
    }
    splice(instData, expInstData);
  }
 
  // order: replace `dim`'s entry with the expanded dim indices in
  // innermost-fastest order; shift every other entry past `dim` up by
  // (expCount - 1).
  SmallVector<int64_t> newOrder;
  DenseI32ArrayAttr orderAttr = getOrder();
  if (orderAttr && !orderAttr.empty()) {
    SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
    newOrder.reserve(newRank);
    for (int64_t o : origOrder) {
      if (o == dim) {
        // Innermost dim of the expanded group is fastest-varying.
        for (int64_t i = expCount - 1; i >= 0; --i)
          newOrder.push_back(dim + i);
      } else if (o > dim) {
        newOrder.push_back(o + expCount - 1);
      } else {
        newOrder.push_back(o);
      }
    }
  }
 
  auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
    if (v.empty())
      return nullptr;
    SmallVector<int32_t> v32(v.begin(), v.end());
    return DenseI32ArrayAttr::get(getContext(), v32);
  };
  return xegpu::LayoutAttr::get(getContext(), toAttr(sgLayout), toAttr(sgData),
                                toAttr(instData), toAttr(laneLayout),
                                toAttr(laneData), toAttr(newOrder));
}
 
// Derive a new layout by transpose the layout using `permutation`.
DistributeLayoutAttr LayoutAttr::transposeDims(ArrayRef<int64_t> permutation) {
 
  SmallVector<int64_t> origSgLayout = getEffectiveSgLayoutAsInt();
  SmallVector<int64_t> origSgData = getEffectiveSgDataAsInt();
  SmallVector<int64_t> origInstData = getEffectiveInstDataAsInt();
  SmallVector<int64_t> origLaneLayout = getEffectiveLaneLayoutAsInt();
  SmallVector<int64_t> origLaneData = getEffectiveLaneDataAsInt();
  SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
 
  SmallVector<int32_t> sgLayout;
  SmallVector<int32_t> sgData;
  SmallVector<int32_t> instData;
  SmallVector<int32_t> laneLayout;
  SmallVector<int32_t> laneData;
  SmallVector<int32_t> order;
 
  for (int64_t idx : permutation) {
    if (!origLaneLayout.empty()) {
      laneLayout.push_back(static_cast<int32_t>(origLaneLayout[idx]));
      laneData.push_back(static_cast<int32_t>(origLaneData[idx]));
    }
    if (!origInstData.empty())
      instData.push_back(static_cast<int32_t>(origInstData[idx]));
    if (!origSgLayout.empty()) {
      sgLayout.push_back(static_cast<int32_t>(origSgLayout[idx]));
      sgData.push_back(static_cast<int32_t>(origSgData[idx]));
    }
  }
 
  // `order` is distinct from the size-valued fields above: its *values* are
  // dimension indices (order[0] is the fastest-varying dim), not per-position
  // sizes. A transpose relabels dimensions (source dim d becomes result dim
  // inversePerm[d]) so the dimension values are remapped through the inverse
  // permutation: newOrder[i] = inversePerm[origOrder[i]].
  //
  // The linearization order this describes is invariant under transpose: a
  // transpose only renames dimensions, so the subgroup ID assigned to a given
  // block of data must stay the same. Remapping the values through the inverse
  // permutation is exactly what preserves that order.
  SmallVector<int64_t> inversePermutation =
      invertPermutationVector(permutation);
  for (int64_t dim : origOrder)
    order.push_back(static_cast<int32_t>(inversePermutation[dim]));
  if (origLaneLayout.empty() && origSgLayout.empty())
    order.clear();
 
  auto toAttr = [&](ArrayRef<int32_t> v) -> DenseI32ArrayAttr {
    return v.empty() ? nullptr : DenseI32ArrayAttr::get(getContext(), v);
  };
  return xegpu::LayoutAttr::get(getContext(), toAttr(sgLayout), toAttr(sgData),
                                toAttr(instData), toAttr(laneLayout),
                                toAttr(laneData), toAttr(order));
}
 
/// Check if this layout is a transpose of another layout.
bool LayoutAttr::isTransposeOf(const xegpu::DistributeLayoutAttr &other,
                               ArrayRef<int64_t> perm,
                               const xegpu::LayoutKind kind) {
  if (!other)
    return false;
  if (getRank() != other.getRank() ||
      perm.size() != static_cast<size_t>(getRank()))
    return false;
  if (!isPermutationVector(perm))
    return false;
  // vector.transpose semantics: dst[i] = src[perm[i]]. So `this` (= dst) is a
  // transpose of `other` (= src) via `perm` iff for all i, dst[i] ==
  // src[perm[i]].
  auto checkTranspose = [](ArrayRef<int64_t> dst, ArrayRef<int64_t> src,
                           ArrayRef<int64_t> perm) {
    for (const auto &ta : llvm::enumerate(perm)) {
      if (dst[ta.index()] != src[ta.value()])
        return false;
    }
    return true;
  };
  // `order` is different: its *values* are dimension indices, so a transpose
  // relabels them through the inverse permutation rather than reindexing by
  // position. `this` (= dst) is a transpose of `other` (= src) iff
  // dst.order[i] == inversePerm[src.order[i]] for all i. This matches the
  // convention produced by `transposeDims`.
  auto checkOrderTranspose = [](ArrayRef<int64_t> dstOrder,
                                ArrayRef<int64_t> srcOrder,
                                ArrayRef<int64_t> perm) {
    if (dstOrder.size() != srcOrder.size())
      return false;
    SmallVector<int64_t> inversePerm = invertPermutationVector(perm);
    for (auto [d, s] : llvm::zip_equal(dstOrder, srcOrder)) {
      if (d != inversePerm[s])
        return false;
    }
    return true;
  };
  if (kind == xegpu::LayoutKind::Subgroup)
    return checkTranspose(getEffectiveSgLayoutAsInt(),
                          other.getEffectiveSgLayoutAsInt(), perm) &&
           checkTranspose(getEffectiveSgDataAsInt(),
                          other.getEffectiveSgDataAsInt(), perm) &&
           checkOrderTranspose(getEffectiveOrderAsInt(),
                               other.getEffectiveOrderAsInt(), perm);
  if (kind == xegpu::LayoutKind::InstData)
    return checkTranspose(getEffectiveInstDataAsInt(),
                          other.getEffectiveInstDataAsInt(), perm);
  if (kind == xegpu::LayoutKind::Lane)
    return checkTranspose(getEffectiveLaneLayoutAsInt(),
                          other.getEffectiveLaneLayoutAsInt(), perm) &&
           checkTranspose(getEffectiveLaneDataAsInt(),
                          other.getEffectiveLaneDataAsInt(), perm) &&
           checkOrderTranspose(getEffectiveOrderAsInt(),
                               other.getEffectiveOrderAsInt(), perm);
 
  return false;
}
 
bool LayoutAttr::isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
                                  SmallVector<int64_t> shape,
                                  xegpu::LayoutKind level) {
  if (!other)
    return false;
  if (getEffectiveOrderAsInt() == other.getEffectiveOrderAsInt()) {
    // short cut when order is the same, no need to compute coords and compare
    if (level == xegpu::LayoutKind::Subgroup)
      if (getEffectiveSgLayoutAsInt() == other.getEffectiveSgLayoutAsInt() &&
          getEffectiveSgDataAsInt() == other.getEffectiveSgDataAsInt())
        return true;
    if (level == xegpu::LayoutKind::Lane)
      if (getEffectiveLaneLayoutAsInt() ==
              other.getEffectiveLaneLayoutAsInt() &&
          getEffectiveLaneDataAsInt() == other.getEffectiveLaneDataAsInt())
        return true;
  }
 
  auto compareCoordsForAllIds = [&](int64_t size) {
    return compareDistributedCoords(*this, other, shape, level, size);
  };
 
  if (level == xegpu::LayoutKind::Subgroup) {
    int64_t wgSize = computeProduct(getEffectiveSgLayoutAsInt());
    return compareCoordsForAllIds(wgSize);
  }
  if (level == xegpu::LayoutKind::InstData) {
    return (getEffectiveInstDataAsInt() == other.getEffectiveInstDataAsInt());
  }
  if (level == xegpu::LayoutKind::Lane) {
    int64_t subgroupSize = computeProduct(getEffectiveLaneLayoutAsInt());
    return compareCoordsForAllIds(subgroupSize);
  }
  return true;
}
 
//===----------------------------------------------------------------------===//
// XeGPU_SliceAttr
//===----------------------------------------------------------------------===//
LogicalResult
SliceAttr::verify(llvm::function_ref<InFlightDiagnostic()> emitError,
                  xegpu::DistributeLayoutAttr parent, DenseI64ArrayAttr dims) {
 
  if (!dims)
    return emitError() << "expected dims attribute";
 
  // check every element in dims is unique and smaller than rank
  llvm::SmallDenseSet<int64_t> seen;
  for (int64_t dim : dims.asArrayRef()) {
    if (dim < 0)
      return emitError() << "invalid dim (" << dim << ") in slice attribute.";
    if (!seen.insert(dim).second)
      return emitError() << "repeated dim (" << dim << ") in slice attribute.";
  }
  return success();
}
 
SliceAttr SliceAttr::flatten() const {
  xegpu::DistributeLayoutAttr parent = getParent();
  SmallVector<DenseI64ArrayAttr> slicedDims({getDims()});
 
  while (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(parent)) {
    parent = sliceAttr.getParent();
    slicedDims.push_back(sliceAttr.getDims());
  }
 
  auto layoutAttr = dyn_cast<xegpu::LayoutAttr>(parent);
  SmallVector<int64_t> indices =
      llvm::to_vector(llvm::seq<int64_t>(0, layoutAttr.getRank()));
 
  // get remaining dims (flattened) by applying slice ops with all slicedDims
  SmallVector<int64_t> remainingDims(indices);
  for (auto dim : llvm::reverse(slicedDims))
    remainingDims = XeGPUDialect::slice(llvm::ArrayRef<int64_t>(remainingDims),
                                        dim.asArrayRef());
 
  // get flattened sliced dims by applying slice ops with the remaining dims
  SmallVector<int64_t> flattenedDims = XeGPUDialect::slice(
      llvm::ArrayRef<int64_t>(indices), llvm::ArrayRef<int64_t>(remainingDims));
 
  return xegpu::SliceAttr::get(
      getContext(), layoutAttr,
      DenseI64ArrayAttr::get(getContext(), flattenedDims));
}
 
FailureOr<SmallVector<Value>>
SliceAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId) {
  SliceAttr attr = flatten();
  auto parent = dyn_cast<LayoutAttr>(attr.getParent());
  return parent.delinearizeId(builder, loc, linearId);
}
 
// Implements DistributeLayoutAttr::computeDistributedCoords to generate
// instructions for computing multi-dimensional offsets when distributed by
// LayoutAttr.
FailureOr<SmallVector<SmallVector<Value>>>
SliceAttr::computeDistributedCoords(OpBuilder &builder, Location loc,
                                    Value linearId, ArrayRef<int64_t> shape) {
  assert(getRank() == static_cast<int64_t>(shape.size()) && "invalid shape.");
 
  SmallVector<int64_t> layout;
  SmallVector<int64_t> subShape;
  if (isForWorkgroup()) {
    layout = getEffectiveSgLayoutAsInt();
    subShape = getEffectiveSgDataAsInt();
  } else if (isForSubgroup()) {
    layout = getEffectiveLaneLayoutAsInt();
    subShape = getEffectiveLaneDataAsInt();
  } else {
    return failure();
  }
 
  if (subShape.empty())
    return failure();
 
  // delinearize Ids
  auto maybeIds = delinearizeId(builder, loc, linearId);
  if (failed(maybeIds))
    return failure();
 
  // The effective sgIds for offsets computing correspond
  // to the dims that are not sliced.
  ArrayRef<int64_t> dims = flatten().getDims().asArrayRef();
  SmallVector<Value> canonicalIds =
      XeGPUDialect::slice(ArrayRef<Value>(*maybeIds), dims);
 
  return genCoordinates(builder, loc, canonicalIds, layout, subShape, shape);
}
 
/// Implements DistributeLayoutAttr::computeStaticDistributedCoords to
/// compute multi-dimensional offsets for a given linear ID when distributed by
/// SliceAttr. Delegates delinearization to the parent LayoutAttr, then uses
/// only the non-sliced dimensions for coordinate computation.
SmallVector<SmallVector<int64_t>>
SliceAttr::computeStaticDistributedCoords(int64_t linearId,
                                          ArrayRef<int64_t> shape) {
  assert(getRank() == static_cast<int64_t>(shape.size()) && "invalid shape.");
 
  SmallVector<int64_t> layout;
  SmallVector<int64_t> subShape;
  SmallVector<int64_t> instData;
  if (isForWorkgroup()) {
    layout = getEffectiveSgLayoutAsInt();
    subShape = getEffectiveSgDataAsInt();
  } else if (isForSubgroup()) {
    instData = getEffectiveInstDataAsInt();
    layout = getEffectiveLaneLayoutAsInt();
    subShape = getEffectiveLaneDataAsInt();
  }
  if (!instData.empty()) {
    linearId = 0;
    subShape = instData;
  }
 
  assert(!subShape.empty() && "sgdata or lanedata cannot be empty");
 
  // Delinearize the ID using the parent layout (same as the IR version).
  SliceAttr flattened = flatten();
  auto parent = dyn_cast<LayoutAttr>(flattened.getParent());
  SmallVector<int64_t> parentLayoutVec;
  if (parent.isForWorkgroup())
    parentLayoutVec = parent.getEffectiveSgLayoutAsInt();
  else
    parentLayoutVec = parent.getEffectiveLaneLayoutAsInt();
 
  SmallVector<int64_t> order = parent.getEffectiveOrderAsInt();
  SmallVector<int64_t> allIds(parentLayoutVec.size());
  int64_t remaining = linearId;
  for (size_t i = 0; i < order.size(); ++i) {
    int64_t dimIdx = order[i];
    allIds[dimIdx] = remaining % parentLayoutVec[dimIdx];
    if (i < order.size() - 1)
      remaining = remaining / parentLayoutVec[dimIdx];
  }
 
  // The effective IDs for coordinate computation correspond
  // to the dims that are not sliced.
  ArrayRef<int64_t> dims = flattened.getDims().asArrayRef();
  SmallVector<int64_t> canonicalIds =
      XeGPUDialect::slice(ArrayRef<int64_t>(allIds), dims);
 
  return genStaticCoordinates(canonicalIds, layout, subShape, shape);
}
 
bool SliceAttr::isSliceOf(const xegpu::DistributeLayoutAttr &other) {
  auto flattenedThis = flatten();
  // If other is a LayoutAttr, just compare directly with parent of
  // flattenedThis.
  if (auto otherLayout = dyn_cast<xegpu::LayoutAttr>(other))
    return flattenedThis.getParent() == otherLayout;
  // If other is a SliceAttr, flatten it first before comparing.
  auto flattenedOther = dyn_cast<xegpu::SliceAttr>(other).flatten();
  // Both must have common parent LayoutAttr.
  if (flattenedThis.getParent() != flattenedOther.getParent())
    return false;
  // otherFlattened's sliced dims must be a subset of flattenedThis's sliced
  // dims.
  llvm::SmallDenseSet<int64_t> thisDims(
      flattenedThis.getDims().asArrayRef().begin(),
      flattenedThis.getDims().asArrayRef().end());
  return llvm::all_of(flattenedOther.getDims().asArrayRef(),
                      [&](int64_t dim) { return thisDims.contains(dim); });
}
 
bool SliceAttr::isEqualTo(const xegpu::DistributeLayoutAttr &other) {
  if (dyn_cast<xegpu::LayoutAttr>(other))
    return false;
 
  auto flattenedThis = flatten();
  auto flattenedOther = dyn_cast<xegpu::SliceAttr>(other).flatten();
 
  return ((flattenedThis.getParent() == flattenedOther.getParent()) &&
          (flattenedThis.getDims() == flattenedOther.getDims()));
}
 
bool SliceAttr::isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
                                 SmallVector<int64_t> shape,
                                 xegpu::LayoutKind level) {
  if (!other)
    return false;
  if (getEffectiveOrderAsInt() == other.getEffectiveOrderAsInt()) {
    // short cut when order is the same, no need to compute coords and compare
    if (level == xegpu::LayoutKind::Subgroup)
      if (getEffectiveSgLayoutAsInt() == other.getEffectiveSgLayoutAsInt() &&
          getEffectiveSgDataAsInt() == other.getEffectiveSgDataAsInt())
        return true;
    if (level == xegpu::LayoutKind::Lane)
      if (getEffectiveLaneLayoutAsInt() ==
              other.getEffectiveLaneLayoutAsInt() &&
          getEffectiveLaneDataAsInt() == other.getEffectiveLaneDataAsInt())
        return true;
  }
 
  auto compareCoordsForAllIds = [&](int64_t size) {
    return compareDistributedCoords(*this, other, shape, level, size);
  };
 
  auto flattenedThis = flatten();
  auto parent = dyn_cast<LayoutAttr>(flattenedThis.getParent());
  if (level == xegpu::LayoutKind::Subgroup) {
    int64_t wgSize = computeProduct(parent.getEffectiveSgLayoutAsInt());
    return compareCoordsForAllIds(wgSize);
  }
  if (level == xegpu::LayoutKind::InstData) {
    return (getEffectiveInstDataAsInt() == other.getEffectiveInstDataAsInt());
  }
  if (level == xegpu::LayoutKind::Lane) {
    int64_t subgroupSize = computeProduct(parent.getEffectiveLaneLayoutAsInt());
    return compareCoordsForAllIds(subgroupSize);
  }
  return true;
}
 
xegpu::SliceAttr SliceAttr::dropSliceDims(ArrayRef<int64_t> sliceDimsToDrop) {
  if (sliceDimsToDrop.empty())
    return *this;
  SmallVector<int64_t> sliceDims{getDims().asArrayRef()};
  for (auto dim : sliceDimsToDrop) {
    auto foundIt = std::find(sliceDims.begin(), sliceDims.end(), dim);
    assert(foundIt != sliceDims.end() &&
           "Expected to find the specified reduction dim in slice dims");
    sliceDims.erase(foundIt);
  }
 
  auto sliceWithoutDims = xegpu::SliceAttr::get(
      this->getContext(), getParent(),
      DenseI64ArrayAttr::get(this->getContext(), sliceDims));
 
  return sliceWithoutDims;
}
 
// Helper function to adjust dimensions from sliced space to parent space
// say we have a parent shape of rank 4, and slice dims [1,3], so the sliced
// shape is of rank 2, if we want to set unit dim [0] in sliced space, it maps
// to dim [0] in parent space; if we want to set unit dim [1] in sliced space,
// it maps to dim [2] in parent space.
static SmallVector<int64_t>
mapSlicedDimsToParentSpace(const SmallVector<int64_t> &dimsToMap,
                           ArrayRef<int64_t> sliceDims) {
  // Rather than recovering the exact parent rank, we compute a safe upper
  // bound so that dimsToMap can be adjusted safely. This upper bound is
  // defined as max(dimsToMap, sliceDims) + 1 + sliceDims.size().
  int64_t maxDim = -1;
  maxDim =
      std::max(maxDim, *std::max_element(sliceDims.begin(), sliceDims.end()));
  maxDim =
      std::max(maxDim, *std::max_element(dimsToMap.begin(), dimsToMap.end()));
  int64_t parentSpaceRank = maxDim + sliceDims.size() + 1;
 
  // get remaining dims in parent space after applying slicing with parent's
  // slice Dims
  llvm::SmallDenseSet<int64_t> slicedDimsSet(sliceDims.begin(),
                                             sliceDims.end());
  SmallVector<int64_t> remainingDims;
  for (int64_t i = 0; i < parentSpaceRank; ++i) {
    if (!slicedDimsSet.contains(i))
      remainingDims.push_back(i);
  }
 
  // Map unit dims from sliced space to parent space
  SmallVector<int64_t> adjustUnitDims;
  for (auto dim : dimsToMap) {
    int64_t mappedDim = remainingDims[dim];
    adjustUnitDims.push_back(mappedDim);
  }
 
  return adjustUnitDims;
}
 
// set the layout for unit dims: sg_data, inst_data and lane_data to 1
DistributeLayoutAttr
SliceAttr::setUnitDimData(SmallVector<int64_t> unitDims) const {
  DistributeLayoutAttr parentLayout = getParent();
 
  ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
 
  SmallVector<int64_t> adjustUnitDims =
      mapSlicedDimsToParentSpace(unitDims, sliceDims);
 
  return SliceAttr::get(getContext(),
                        parentLayout.setUnitDimData(adjustUnitDims), getDims());
}
 
// set the layout for the sepcified unit dims: sg_lane and lane_layout to 1
DistributeLayoutAttr
SliceAttr::setUnitDimLayout(SmallVector<int64_t> unitDims) const {
  DistributeLayoutAttr parentLayout = getParent();
 
  ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
 
  SmallVector<int64_t> adjustUnitDims =
      mapSlicedDimsToParentSpace(unitDims, sliceDims);
 
  return SliceAttr::get(
      getContext(), parentLayout.setUnitDimLayout(adjustUnitDims), getDims());
}
 
// Derive a new layout with sg_data, inst_data and lane_data set to the
// specified values for the given dimension
DistributeLayoutAttr SliceAttr::setDimData(int64_t dim, int64_t sgData,
                                           int64_t instData, int64_t laneData) {
  ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
  auto parent = getParent();
 
  SmallVector<int64_t> dimSet;
  dimSet.push_back(dim);
  SmallVector<int64_t> adjustDims =
      mapSlicedDimsToParentSpace(dimSet, sliceDims);
  return SliceAttr::get(
      getContext(),
      parent.setDimData(adjustDims[0], sgData, instData, laneData), getDims());
}
 
// Derive a new layout by removing dimensions. `dimGroup` specifies a group of
// dimensions to be removed in the derived layout.
//
// Example: drop the 2nd dimension from a rank-3 sliced view.
//
// Suppose:
//   xegpu.layout = slice<layout<[V0, V1, V2, V3, V4]>, [1, 3]>
//
// The slice removes parent dims [1, 3], so the sliced-space dims map to
// parent dims [V0, V2, V4].
//
// If we drop sliced-space dim 1 (the 2nd dim), that corresponds to dropping
// parent dim 2, result in parent layout [V0, V1, V3, V4] after dropping.
// After parent dim 2 is removed, sliced dims [1, 3] must be reindexed to [1,
// 2].
//
// Result:
//   xegpu.layout = slice<layout<[0, 1, 3, 4]>, [1, 2]>
DistributeLayoutAttr SliceAttr::dropDims(SmallVector<int64_t> dimGroup) {
  // Map the sliced dims from parent space to collapsed space
  SmallVector<int64_t> sliceDims = llvm::to_vector(getDims().asArrayRef());
  SmallVector<int64_t> dimsInParentSpace =
      mapSlicedDimsToParentSpace(dimGroup, sliceDims);
 
  auto droppedParent = getParent().dropDims(dimsInParentSpace);
 
  // Adjust the sliced dims after dropping dims in parent space. For example, if
  // we drop dim 2 in parent space, the dims after dim 2 will all be shifted by
  // 1, so sliced dim 3 will be adjusted to 2.
  SmallVector<int64_t> newSliceDims;
  for (int64_t d : sliceDims) {
    int64_t offset =
        llvm::count_if(dimsInParentSpace, [&](int64_t s) { return s < d; });
    newSliceDims.push_back(d - offset);
  }
 
  return SliceAttr::get(getContext(), droppedParent,
                        DenseI64ArrayAttr::get(getContext(), newSliceDims));
}
 
// Derive a new layout by collapsing dimensions.
// `dimGroup` specifies a group of adjacent dimensions
// that are collapsed into a single dimension in the derived layout.
DistributeLayoutAttr SliceAttr::collapseDims(SmallVector<int64_t> dimGroup) {
 
  // Map the sliced dims from parent space to collapsed space
  SmallVector<int64_t> sliceDims = llvm::to_vector(getDims().asArrayRef());
  assert("expect sliceDims not being collapsed" &&
         llvm::none_of(dimGroup, [&](int64_t dim) {
           return llvm::is_contained(sliceDims, dim);
         }));
  SmallVector<int64_t> dimsInParentSpace =
      mapSlicedDimsToParentSpace(dimGroup, sliceDims);
 
  auto collapsedParent = getParent().collapseDims(dimsInParentSpace);
  return SliceAttr::get(getContext(), collapsedParent,
                        DenseI64ArrayAttr::get(getContext(), sliceDims));
}
 
// Derive a new layout by expanding a single sliced-space dim into multiple
// adjacent dims. The dim is mapped to parent space, the parent layout is
// expanded there, and the slice dims that lie past the expanded position
// are shifted up by `targetShape.size() - 1`.
DistributeLayoutAttr SliceAttr::expandDim(int64_t dim,
                                          ArrayRef<int64_t> targetShape) {
  // `dim` is in slice space; map it to parent space (parent dims listed in
  // `sliceDims` are removed by the slice, so the mapping always lands on a
  // non-sliced parent dim).
  ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
  SmallVector<int64_t> dimSet = {dim};
  SmallVector<int64_t> dimsInParentSpace =
      mapSlicedDimsToParentSpace(dimSet, sliceDims);
  int64_t parentDim = dimsInParentSpace[0];
 
  auto expandedParent = getParent().expandDim(parentDim, targetShape);
 
  int64_t shift = static_cast<int64_t>(targetShape.size()) - 1;
  SmallVector<int64_t> newSliceDims;
  newSliceDims.reserve(sliceDims.size());
  for (int64_t s : sliceDims)
    newSliceDims.push_back(s > parentDim ? s + shift : s);
 
  return SliceAttr::get(getContext(), expandedParent,
                        DenseI64ArrayAttr::get(getContext(), newSliceDims));
}
 
SmallVector<int64_t> getPermForParentLayout(ArrayRef<int64_t> sliceDims,
                                            ArrayRef<int64_t> permutation) {
  SmallVector<int64_t> sortedSliceDims = llvm::to_vector(sliceDims);
  llvm::sort(sortedSliceDims);
 
  for (size_t i = 1; i < sortedSliceDims.size(); ++i) {
    assert((sortedSliceDims[i] == sortedSliceDims[i - 1] + 1) &&
           "slice dims non consecutive, cannot be transposed");
  }
 
  SmallVector<int64_t> permForParent;
  if (sortedSliceDims.front() == 0) {
    // Example: sliceDims.size() = 2, permutation= {1, 0}
    // result: {3, 2, 1, 0}.
    for (int64_t dim : permutation)
      permForParent.push_back(dim + sortedSliceDims.size());
    for (int64_t i = sortedSliceDims.size() - 1; i >= 0; --i)
      permForParent.push_back(i);
  } else {
    // Example: sliceDims.size() = 2, permutation = {0, 1}
    // result: {3, 2, 0, 1}.
    for (int64_t i = sortedSliceDims.size() - 1; i >= 0; --i)
      permForParent.push_back(i + permutation.size());
    for (int64_t dim : permutation)
      permForParent.push_back(dim);
  }
  return permForParent;
}
 
// Derive a new layout by transpose the layout using `permutation`.
DistributeLayoutAttr SliceAttr::transposeDims(ArrayRef<int64_t> permutation) {
  SmallVector<int64_t> sliceDims = llvm::to_vector(getDims().asArrayRef());
  DistributeLayoutAttr parent = getParent();
  SmallVector<int64_t> permForParent =
      getPermForParentLayout(sliceDims, permutation);
  auto transposedParent = parent.transposeDims(permForParent);
  return SliceAttr::get(getContext(), transposedParent,
                        DenseI64ArrayAttr::get(getContext(), sliceDims));
}
 
/// Check if this layout is a transpose of another layout.
bool SliceAttr::isTransposeOf(const xegpu::DistributeLayoutAttr &other,
                              ArrayRef<int64_t> perm,
                              const xegpu::LayoutKind kind) {
  // other must be a SliceAttr with the same slice dims.
  auto otherSlice = dyn_cast<xegpu::SliceAttr>(other);
  if (!otherSlice || getDims() != otherSlice.getDims())
    return false;
  // check whether the parent layout is transpose of each other.
  SmallVector<int64_t> sliceDims = llvm::to_vector(getDims().asArrayRef());
  DistributeLayoutAttr parent = getParent();
  SmallVector<int64_t> permForParent = getPermForParentLayout(sliceDims, perm);
  auto otherParent = otherSlice.getParent();
  return parent.isTransposeOf(otherParent, permForParent, kind);
}
 
//===----------------------------------------------------------------------===//
// XeGPU_RangeAttr
//===----------------------------------------------------------------------===//
 
LogicalResult
RangeAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
                  IntegerAttr startOfRange, IntegerAttr endOfRange) {
  if (startOfRange.getInt() >= endOfRange.getInt())
    return emitError() << "'end' : " << endOfRange.getInt()
                       << " must be greater than 'start' : "
                       << startOfRange.getInt();
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_TensorDescType
//===----------------------------------------------------------------------===//
 
mlir::Type TensorDescType::parse(AsmParser &parser) {
  llvm::SmallVector<int64_t> shape;
  mlir::Type elementType;
  mlir::FailureOr<mlir::Attribute> encoding;
  mlir::FailureOr<mlir::Attribute> layout;
 
  // Parse literal '<'
  if (parser.parseLess())
    return {};
 
  auto shapeLoc = parser.getCurrentLocation();
  if (mlir::failed(parser.parseDimensionList(shape))) {
    parser.emitError(shapeLoc, "failed to parse parameter 'shape'");
    return {};
  }
 
  auto elemTypeLoc = parser.getCurrentLocation();
  if (mlir::failed(parser.parseType(elementType))) {
    parser.emitError(elemTypeLoc, "failed to parse parameter 'elementType'");
    return {};
  }
 
  // parse optional attributes
  while (mlir::succeeded(parser.parseOptionalComma())) {
    mlir::Attribute attr;
    ParseResult res = parser.parseAttribute(attr);
    if (mlir::succeeded(res)) {
      if (mlir::isa<DistributeLayoutAttr>(attr)) {
        layout = attr;
        continue;
      }
      if (mlir::isa<BlockTensorDescAttr>(attr)) {
        encoding = attr;
        continue;
      }
    }
    return {};
  }
 
  // Parse literal '>'
  if (parser.parseGreater())
    return {};
 
  MLIRContext *ctxt = parser.getContext();
  return TensorDescType::getChecked(
      [&]() { return parser.emitError(parser.getNameLoc()); }, ctxt, shape,
      elementType, encoding.value_or(BlockTensorDescAttr::get(ctxt)),
      layout.value_or(mlir::Attribute()));
}
 
void TensorDescType::print(AsmPrinter &printer) const {
  printer << "<";
 
  auto shape = getShape();
  for (int64_t dim : shape) {
    if (mlir::ShapedType::isDynamic(dim))
      printer << '?';
    else
      printer << dim;
    printer << 'x';
  }
 
  printer << getElementType();
 
  auto encoding = getEncoding();
  auto blockAttr = llvm::dyn_cast_if_present<BlockTensorDescAttr>(encoding);
  if (encoding && (!blockAttr || !blockAttr.hasDefaultsOnly()))
    printer << ", " << encoding;
 
  if (auto layout = getLayout())
    printer << ", " << layout;
 
  printer << ">";
}
 
TensorDescType TensorDescType::get(llvm::ArrayRef<int64_t> shape,
                                   mlir::Type elementType, int array_length,
                                   bool boundary_check,
                                   MemorySpace memory_space,
                                   mlir::Attribute layout) {
  auto *context = elementType.getContext();
  auto attr = BlockTensorDescAttr::get(context, memory_space, array_length,
                                       boundary_check);
  return Base::get(context, shape, elementType, attr, layout);
}
 
LogicalResult
TensorDescType::verify(llvm::function_ref<InFlightDiagnostic()> emitError,
                       llvm::ArrayRef<int64_t> shape, mlir::Type elementType,
                       mlir::Attribute encoding, mlir::Attribute layout) {
  size_t rank = shape.size();
 
  if (rank == 0)
    return emitError() << "expected non-zero rank tensor";
 
  auto blockAttr = mlir::dyn_cast_if_present<BlockTensorDescAttr>(encoding);
  if (blockAttr) {
    MemorySpaceAttr memorySpaceAttr = blockAttr.getMemorySpace();
    if (rank > 1 && memorySpaceAttr &&
        memorySpaceAttr.getValue() == MemorySpace::SLM)
      return emitError() << "SLM is only supported for 1D block tensor";
  }
 
  if (!elementType.isIntOrFloat())
    return emitError() << "unsupported element type " << elementType
                       << ": expected integer or float";
 
  if (auto layoutAttr =
          mlir::dyn_cast_if_present<DistributeLayoutAttr>(layout)) {
    if (rank != (size_t)layoutAttr.getRank())
      return emitError() << "expected layout rank to match tensor rank";
 
    if (!layoutAttr.isDistributable(SmallVector<int64_t>(shape))) {
      std::string shapeStr;
      llvm::raw_string_ostream stream(shapeStr);
      llvm::interleaveComma(shape, stream);
      return emitError() << "cannot distribute [" << shapeStr << "] using "
                         << layoutAttr;
    }
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// XeGPU_MemDescType
//===----------------------------------------------------------------------===//
mlir::Type MemDescType::parse(AsmParser &parser) {
  llvm::SmallVector<int64_t> shape;
  mlir::Type elementType;
  mlir::FailureOr<MemLayoutAttr> layout;
 
  // Parse literal '<'
  if (parser.parseLess())
    return {};
 
  auto shapeLoc = parser.getCurrentLocation();
  if (mlir::failed(parser.parseDimensionList(shape, false, true))) {
    parser.emitError(shapeLoc, "failed to parse parameter 'shape'");
    return {};
  }
 
  auto elemTypeLoc = parser.getCurrentLocation();
  if (mlir::failed(parser.parseType(elementType))) {
    parser.emitError(elemTypeLoc, "failed to parse parameter 'elementType'");
    return {};
  }
 
  // parse optional attributes
  if (mlir::succeeded(parser.parseOptionalComma())) {
    MemLayoutAttr attr;
    ParseResult res = parser.parseAttribute(attr);
    if (mlir::failed(res))
      return {};
    layout = attr;
  }
 
  // Parse literal '>'
  if (parser.parseGreater())
    return {};
 
  MLIRContext *ctxt = parser.getContext();
  return MemDescType::getChecked(
      [&]() { return parser.emitError(parser.getNameLoc()); }, ctxt, shape,
      elementType, layout.value_or(MemLayoutAttr()));
}
 
void MemDescType::print(AsmPrinter &printer) const {
  printer << "<";
 
  printer.printDimensionList(getShape());
  printer << 'x';
  printer << getElementType();
 
  if (auto layout = getMemLayout())
    printer << ", " << layout;
 
  printer << ">";
}
 
//===----------------------------------------------------------------------===//
// XeGPU_MemDescType
//===----------------------------------------------------------------------===//
 
Attribute MemLayoutAttr::parse(AsmParser &parser, Type type) {
 
  auto *context = parser.getContext();
  llvm::SMLoc loc = parser.getCurrentLocation();
 
  llvm::SmallDenseSet<StringRef> seenKeys;
  SmallVector<NamedAttribute> attributes;
 
  auto parseElt = [&]() -> ParseResult {
    StringRef nameId;
    if (failed(parser.parseKeyword(&nameId)))
      return parser.emitError(loc, "expected valid attribute name");
 
    if (!seenKeys.insert(nameId).second)
      return parser.emitError(loc, "duplicate key '")
             << nameId << " in mem layout attribute";
 
    if (failed(parser.parseEqual()))
      return failure();
 
    Attribute attr;
    if (failed(parser.parseAttribute(attr)))
      return failure();
    attributes.emplace_back(nameId, attr);
    return success();
  };
 
  // Parse literal '<'
  if (parser.parseLess())
    return {};
 
  if (failed(parser.parseCommaSeparatedList(parseElt)))
    return {};
 
  // Parse literal '>'
  if (parser.parseGreater())
    return {};
 
  return parser.getChecked<MemLayoutAttr>(
      loc, context, DictionaryAttr::get(context, attributes));
}
 
void MemLayoutAttr::print(AsmPrinter &printer) const {
  printer << "<";
  ArrayRef<NamedAttribute> attrs = getAttrs().getValue();
  for (size_t i = 0; i < attrs.size(); i++) {
    printer << attrs[i].getName().str() << " = " << attrs[i].getValue();
    if (i < attrs.size() - 1)
      printer << ", ";
  }
  printer << ">";
}
// a helper utility to perform binary operation on OpFoldResult.
// If both a and b are attributes, it will simply return the result.
// Otherwise, the corresponding arith op will be generated, and an
// contant op will be created if one of them is an attribute.
template <typename ArithOp>
OpFoldResult genBinOp(OpFoldResult a, OpFoldResult b, Location loc,
                      OpBuilder &builder) {
  auto aVal = getValueOrCreateConstantIndexOp(builder, loc, a);
  auto bVal = getValueOrCreateConstantIndexOp(builder, loc, b);
  return ArithOp::create(builder, loc, aVal, bVal).getResult();
}
 
// a helper utility to perform division operation on OpFoldResult and int64_t.
#define div(a, b)                                                              \
  genBinOp<arith::DivSIOp>(a, builder.getIndexAttr(b), loc, builder)
 
// a helper utility to perform reminder operation on OpFoldResult and int64_t.
#define rem(a, b)                                                              \
  genBinOp<arith::RemSIOp>(a, builder.getIndexAttr(b), loc, builder)
 
// a helper utility to perform multiply operation on OpFoldResult and int64_t.
#define mul(a, b)                                                              \
  genBinOp<arith::MulIOp>(a, builder.getIndexAttr(b), loc, builder)
 
// a helper utility to perform addition operation on two OpFoldResult.
#define add(a, b) genBinOp<arith::AddIOp>(a, b, loc, builder)
 
// block the given offsets according to the block shape
// say the original offset is [y, x], and the block shape is [By, Bx],
// then the blocked offset is [y/By, x/Bx, y%By, x%Bx]
SmallVector<OpFoldResult> getBlockedOffsets(OpBuilder &builder, Location loc,
                                            ArrayRef<OpFoldResult> offsets,
                                            ArrayRef<int64_t> blockShape) {
 
  assert(offsets.size() == blockShape.size() &&
         "offsets and blockShape must have the same size");
  SmallVector<OpFoldResult> blockedOffsets;
  SmallVector<OpFoldResult> divs, rems;
 
  for (auto [offset, block] : llvm::zip(offsets, blockShape)) {
    divs.push_back(div(offset, block));
    rems.push_back(rem(offset, block));
  }
  blockedOffsets.append(divs.begin(), divs.end());
  blockedOffsets.append(rems.begin(), rems.end());
 
  return blockedOffsets;
}
 
// Get strides as vector of integer for MemDesc.
SmallVector<int64_t> MemDescType::getStrideShape() {
 
  SmallVector<int64_t> matrixShape(getShape().begin(), getShape().end());
 
  ArrayAttr strideAttr = getStrideAttr();
  SmallVector<int64_t> strides;
  for (Attribute attr : strideAttr.getValue()) {
    strides.push_back(cast<IntegerAttr>(attr).getInt());
  }
 
  SmallVector<int64_t> innerBlkShape = getBlockShape();
 
  // get perm from FCD to LCD
  // perm[i] = the dim with i-th smallest stride
  SmallVector<int, 4> perm =
      llvm::to_vector<4>(llvm::seq<int>(0, strides.size()));
  llvm::sort(perm, [&](int a, int b) { return strides[a] < strides[b]; });
 
  assert(strides[perm[0]] == 1 && "inner most dim must have stride 1");
 
  SmallVector<int64_t> innerBlkStride(innerBlkShape.size());
  innerBlkStride[perm[0]] = 1;
  for (size_t i = 1; i < perm.size(); ++i)
    innerBlkStride[perm[i]] =
        innerBlkStride[perm[i - 1]] * innerBlkShape[perm[i - 1]];
 
  // compute the original matrix shape using the stride info
  // and compute the number of blocks in each dimension
  // The shape of highest dim can't be derived from stride info,
  // but doesn't impact the stride computation for blocked layout.
  SmallVector<int64_t> matrixShapeOrig(matrixShape.size());
  SmallVector<int64_t> BlkShapeOrig(matrixShape.size());
  for (size_t i = 0; i < perm.size() - 1; ++i) {
    matrixShapeOrig[perm[i]] = strides[perm[i + 1]] / strides[perm[i]];
    BlkShapeOrig[perm[i]] = matrixShapeOrig[perm[i]] / innerBlkShape[perm[i]];
  }
 
  int64_t innerBlkSize = 1;
  for (auto s : innerBlkShape)
    innerBlkSize *= s;
 
  SmallVector<int64_t> outerBlkStride(matrixShape.size());
  outerBlkStride[perm[0]] = innerBlkSize;
  for (size_t i = 0; i < perm.size() - 1; ++i) {
    outerBlkStride[perm[i + 1]] =
        outerBlkStride[perm[i]] * BlkShapeOrig[perm[i]];
  }
 
  // combine the inner and outer strides
  SmallVector<int64_t> blockedStrides;
  blockedStrides.append(outerBlkStride.begin(), outerBlkStride.end());
  blockedStrides.append(innerBlkStride.begin(), innerBlkStride.end());
 
  return blockedStrides;
}
 
// Calculate the linear offset using the blocked offsets and stride
Value MemDescType::getLinearOffsets(OpBuilder &builder, Location loc,
                                    ArrayRef<OpFoldResult> offsets) {
 
  SmallVector<int64_t> matrixShape(getShape().begin(), getShape().end());
  SmallVector<int64_t> blockShape = getBlockShape();
  SmallVector<int64_t> strides = getStrideShape();
  SmallVector<OpFoldResult> blockedOffsets;
 
  // blockshape equal to matrixshape means no blocking
  if (llvm::equal(blockShape, matrixShape)) {
    // remove the outer dims from strides
    strides.erase(strides.begin(), strides.begin() + matrixShape.size());
  } else {
    assert(offsets.size() == blockShape.size() &&
           "offsets and blockShape must have the same size");
    // say the original offset is [y, x], and the block shape is [By, Bx],
    // then the blocked offset is [y/By, x/Bx, y%By, x%Bx]
 
    SmallVector<OpFoldResult> divs, rems;
 
    for (auto [offset, block] : llvm::zip(offsets, blockShape)) {
      divs.push_back(div(offset, block));
      rems.push_back(rem(offset, block));
    }
    blockedOffsets.append(divs.begin(), divs.end());
    blockedOffsets.append(rems.begin(), rems.end());
    offsets = blockedOffsets;
  }
 
  // Start with initial value as matrix descriptor's base offset.
  Value linearOffset = arith::ConstantIndexOp::create(builder, loc, 0);
  for (size_t i = 0; i < offsets.size(); ++i) {
    OpFoldResult mulResult = mul(offsets[i], strides[i]);
    Value mulVal = getValueOrCreateConstantIndexOp(builder, loc, mulResult);
    linearOffset = arith::AddIOp::create(builder, loc, mulVal, linearOffset);
  }
 
  return linearOffset;
}
 
} // namespace xegpu
} // namespace mlir
 
#include <mlir/Dialect/XeGPU/IR/XeGPUDialect.cpp.inc>
#define GET_ATTRDEF_CLASSES
#include <mlir/Dialect/XeGPU/IR/XeGPUAttrs.cpp.inc>
#define GET_TYPEDEF_CLASSES
#include <mlir/Dialect/XeGPU/IR/XeGPUTypes.cpp.inc>