/*===- TableGen'erated file -------------------------------------*- C++ -*-===*\ |* *| |* Interface Declarations *| |* *| |* Automatically generated file, do not edit! *| |* *| \*===----------------------------------------------------------------------===*/ namespace mlir::xegpu { /// Common trait for all XeGPU layouts. class DistributeLayoutAttr; } // namespace mlir::xegpu namespace mlir::xegpu { /// Common trait for all XeGPU layouts. namespace detail { struct DistributeLayoutAttrInterfaceTraits { struct Concept { /// The methods defined by the interface. bool (*isForWorkgroup)(const Concept *impl, ::mlir::Attribute ); bool (*isForSubgroup)(const Concept *impl, ::mlir::Attribute ); bool (*isForLane)(const Concept *impl, ::mlir::Attribute ); int64_t (*getRank)(const Concept *impl, ::mlir::Attribute ); int64_t (*getNumSubgroups)(const Concept *impl, ::mlir::Attribute ); DenseI32ArrayAttr (*getOrder)(const Concept *impl, ::mlir::Attribute ); SmallVector (*getEffectiveSgLayoutAsInt)(const Concept *impl, ::mlir::Attribute ); SmallVector (*getEffectiveSgDataAsInt)(const Concept *impl, ::mlir::Attribute ); SmallVector (*getEffectiveInstDataAsInt)(const Concept *impl, ::mlir::Attribute ); SmallVector (*getEffectiveLaneLayoutAsInt)(const Concept *impl, ::mlir::Attribute ); SmallVector (*getEffectiveLaneDataAsInt)(const Concept *impl, ::mlir::Attribute ); SmallVector (*getEffectiveOrderAsInt)(const Concept *impl, ::mlir::Attribute ); xegpu::DistributeLayoutAttr (*dropSgLayoutAndData)(const Concept *impl, ::mlir::Attribute ); xegpu::DistributeLayoutAttr (*dropInstData)(const Concept *impl, ::mlir::Attribute ); xegpu::DistributeLayoutAttr (*setUnitDimData)(const Concept *impl, ::mlir::Attribute , const SmallVector); xegpu::DistributeLayoutAttr (*setUnitDimLayout)(const Concept *impl, ::mlir::Attribute , const SmallVector); FailureOr> (*delinearizeId)(const Concept *impl, ::mlir::Attribute , OpBuilder &, Location, Value); xegpu::DistributeLayoutAttr (*setDimData)(const Concept *impl, ::mlir::Attribute , int64_t, int64_t, int64_t, int64_t); xegpu::DistributeLayoutAttr (*dropDims)(const Concept *impl, ::mlir::Attribute , SmallVector); xegpu::DistributeLayoutAttr (*collapseDims)(const Concept *impl, ::mlir::Attribute , SmallVector); xegpu::DistributeLayoutAttr (*expandDim)(const Concept *impl, ::mlir::Attribute , int64_t, ArrayRef); xegpu::DistributeLayoutAttr (*transposeDims)(const Concept *impl, ::mlir::Attribute , ArrayRef); FailureOr>> (*computeDistributedCoords)(const Concept *impl, ::mlir::Attribute , OpBuilder &, Location, Value, ArrayRef); SmallVector> (*computeStaticDistributedCoords)(const Concept *impl, ::mlir::Attribute , int64_t, ArrayRef); FailureOr> (*computeDistributedShape)(const Concept *impl, ::mlir::Attribute , SmallVector); bool (*isDistributable)(const Concept *impl, ::mlir::Attribute , SmallVector); bool (*isSliceOf)(const Concept *impl, ::mlir::Attribute , const xegpu::DistributeLayoutAttr&); bool (*isTransposeOf)(const Concept *impl, ::mlir::Attribute , const xegpu::DistributeLayoutAttr&, ArrayRef, xegpu::LayoutKind); bool (*isCompatibleWith)(const Concept *impl, ::mlir::Attribute , const xegpu::DistributeLayoutAttr&, SmallVector, xegpu::LayoutKind); bool (*isEqualTo)(const Concept *impl, ::mlir::Attribute , const xegpu::DistributeLayoutAttr&); }; template class Model : public Concept { public: using Interface = ::mlir::xegpu::DistributeLayoutAttr; Model() : Concept{isForWorkgroup, isForSubgroup, isForLane, getRank, getNumSubgroups, getOrder, getEffectiveSgLayoutAsInt, getEffectiveSgDataAsInt, getEffectiveInstDataAsInt, getEffectiveLaneLayoutAsInt, getEffectiveLaneDataAsInt, getEffectiveOrderAsInt, dropSgLayoutAndData, dropInstData, setUnitDimData, setUnitDimLayout, delinearizeId, setDimData, dropDims, collapseDims, expandDim, transposeDims, computeDistributedCoords, computeStaticDistributedCoords, computeDistributedShape, isDistributable, isSliceOf, isTransposeOf, isCompatibleWith, isEqualTo} {} static inline bool isForWorkgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline bool isForSubgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline bool isForLane(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline int64_t getRank(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline int64_t getNumSubgroups(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline DenseI32ArrayAttr getOrder(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveSgLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveSgDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveInstDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveLaneLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveLaneDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveOrderAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline xegpu::DistributeLayoutAttr dropSgLayoutAndData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline xegpu::DistributeLayoutAttr dropInstData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline xegpu::DistributeLayoutAttr setUnitDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims); static inline xegpu::DistributeLayoutAttr setUnitDimLayout(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims); static inline FailureOr> delinearizeId(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId); static inline xegpu::DistributeLayoutAttr setDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, int64_t sgData, int64_t instData, int64_t laneData); static inline xegpu::DistributeLayoutAttr dropDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup); static inline xegpu::DistributeLayoutAttr collapseDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup); static inline xegpu::DistributeLayoutAttr expandDim(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, ArrayRef targetShape); static inline xegpu::DistributeLayoutAttr transposeDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, ArrayRef permutation); static inline FailureOr>> computeDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId, ArrayRef shape); static inline SmallVector> computeStaticDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t linearId, ArrayRef shape); static inline FailureOr> computeDistributedShape(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape); static inline bool isDistributable(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape); static inline bool isSliceOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other); static inline bool isTransposeOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, ArrayRef perm, xegpu::LayoutKind kind); static inline bool isCompatibleWith(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, SmallVector shape, xegpu::LayoutKind level); static inline bool isEqualTo(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other); }; template class FallbackModel : public Concept { public: using Interface = ::mlir::xegpu::DistributeLayoutAttr; FallbackModel() : Concept{isForWorkgroup, isForSubgroup, isForLane, getRank, getNumSubgroups, getOrder, getEffectiveSgLayoutAsInt, getEffectiveSgDataAsInt, getEffectiveInstDataAsInt, getEffectiveLaneLayoutAsInt, getEffectiveLaneDataAsInt, getEffectiveOrderAsInt, dropSgLayoutAndData, dropInstData, setUnitDimData, setUnitDimLayout, delinearizeId, setDimData, dropDims, collapseDims, expandDim, transposeDims, computeDistributedCoords, computeStaticDistributedCoords, computeDistributedShape, isDistributable, isSliceOf, isTransposeOf, isCompatibleWith, isEqualTo} {} static inline bool isForWorkgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline bool isForSubgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline bool isForLane(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline int64_t getRank(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline int64_t getNumSubgroups(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline DenseI32ArrayAttr getOrder(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveSgLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveSgDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveInstDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveLaneLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveLaneDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline SmallVector getEffectiveOrderAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline xegpu::DistributeLayoutAttr dropSgLayoutAndData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline xegpu::DistributeLayoutAttr dropInstData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val); static inline xegpu::DistributeLayoutAttr setUnitDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims); static inline xegpu::DistributeLayoutAttr setUnitDimLayout(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims); static inline FailureOr> delinearizeId(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId); static inline xegpu::DistributeLayoutAttr setDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, int64_t sgData, int64_t instData, int64_t laneData); static inline xegpu::DistributeLayoutAttr dropDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup); static inline xegpu::DistributeLayoutAttr collapseDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup); static inline xegpu::DistributeLayoutAttr expandDim(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, ArrayRef targetShape); static inline xegpu::DistributeLayoutAttr transposeDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, ArrayRef permutation); static inline FailureOr>> computeDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId, ArrayRef shape); static inline SmallVector> computeStaticDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t linearId, ArrayRef shape); static inline FailureOr> computeDistributedShape(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape); static inline bool isDistributable(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape); static inline bool isSliceOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other); static inline bool isTransposeOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, ArrayRef perm, xegpu::LayoutKind kind); static inline bool isCompatibleWith(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, SmallVector shape, xegpu::LayoutKind level); static inline bool isEqualTo(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other); }; template class ExternalModel : public FallbackModel { public: using ConcreteEntity = ConcreteAttr; }; }; template struct DistributeLayoutAttrTrait; } // namespace detail class DistributeLayoutAttr : public ::mlir::AttributeInterface { public: using ::mlir::AttributeInterface::AttributeInterface; template struct Trait : public detail::DistributeLayoutAttrTrait {}; /// Check the availability of workgroup level layouts bool isForWorkgroup() const; /// Check the availability of subgroup level layouts bool isForSubgroup() const; /// Check the availability of lane level layouts bool isForLane() const; /// Get the rank of attribute int64_t getRank() const; /// Get the num of effective subgroups int64_t getNumSubgroups() const; /// Get the order of the layout attribute DenseI32ArrayAttr getOrder() const; /// Get the effective SgLayout of the layout attribute as integer array SmallVector getEffectiveSgLayoutAsInt() const; /// Get the effective SgData of the layout attribute as integer array SmallVector getEffectiveSgDataAsInt() const; /// Get the effective InstData of the layout attribute as integer array SmallVector getEffectiveInstDataAsInt() const; /// Get the effective LaneLayout of the layout attribute as integer array SmallVector getEffectiveLaneLayoutAsInt() const; /// Get the effective LaneData of the layout attribute as integer array SmallVector getEffectiveLaneDataAsInt() const; /// Get the effective order of the layout attribute as integer array SmallVector getEffectiveOrderAsInt() const; /// Derive a new layout by dropping sgLayout and sgData xegpu::DistributeLayoutAttr dropSgLayoutAndData() const; /// Derive a new layout by dropping InstData xegpu::DistributeLayoutAttr dropInstData() const; /// Derive a new layout with sg_data, inst_data and lane_data set to 1 for the specified unit dims xegpu::DistributeLayoutAttr setUnitDimData(const SmallVector unitDims) const; /// Derive a new layout with sg_lane and lane_layout set to 1 for the specified unit dims xegpu::DistributeLayoutAttr setUnitDimLayout(const SmallVector unitDims) const; /// Delinearizes a linear ID into its multidimensional /// indices based on the effective layout level. FailureOr> delinearizeId(OpBuilder & builder, Location loc, Value linearId) const; /// Derive a new layout with sg_data, inst_data and lane_data set to the /// specified values for the given dimension. Passing -1 for any parameter /// preserves its original value. xegpu::DistributeLayoutAttr setDimData(int64_t dim, int64_t sgData, int64_t instData, int64_t laneData) const; /// Derive a new layout by removing dimensions. /// `dimGroup` specifies a group of dimensions to be removed in the derived layout. xegpu::DistributeLayoutAttr dropDims(SmallVector dimGroup) const; /// 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. xegpu::DistributeLayoutAttr collapseDims(SmallVector dimGroup) const; /// Derive a new layout by expanding a single dimension into /// multiple adjacent dimensions whose extents are given by `targetShape` /// (`product(targetShape) == originalSize` on `dim`). The original layout's /// per-dim values on `dim` are distributed across the new dims: /// `sg_layout` and `lane_layout` spread outer-to-inner (capped per dim by /// `targetShape[i]`); `sg_data`, `lane_data`, and `inst_data` fill innermost- /// first; `inst_data` is seeded from `lane_layout * lane_data` per new dim. /// `order`, when present, is rewritten so the expanded dims appear /// innermost-fastest in the original dim's slot; if the input has no /// `order`, the result has no `order` either (the default FCD-first /// order `[rank-1, ..., 0]` is assumed). /// Example: layout with /// expandDim(1, [2, 8]) -> layout /// (no order in, no order out). xegpu::DistributeLayoutAttr expandDim(int64_t dim, ArrayRef targetShape) const; /// Derive a new layout by trasnposing it using `permutation`. xegpu::DistributeLayoutAttr transposeDims(ArrayRef permutation) const; /// Generates instructions to compute multidimensional coordinates for dist units /// assigned to a level identified by linearId. The shape parameter /// represents the higher-level problem size. Each level may access /// multiple blocks according to round-robin distribution rules. FailureOr>> computeDistributedCoords(OpBuilder & builder, Location loc, Value linearId, ArrayRef shape) const; /// Statically computes multidimensional coordinates for all dist units /// assigned to a compute unit identified by `linearId`. This is the /// compile-time counterpart of `computeDistributedCoords`: it performs /// the same delinearization and round-robin enumeration but operates /// entirely on static integer values. Returns a list of coordinate /// vectors, one per dist unit. SmallVector> computeStaticDistributedCoords(int64_t linearId, ArrayRef shape) const; /// Computes the distributed shape for each compute unit by dividing each /// dimension of `shape` by the corresponding layout factor (sg_layout or /// lane_layout). /// The distributed shape represents the per-compute-unit tile. Each /// distribution unit is defined as the combination of layout factors and /// per-unit data (`subshape`, e.g., sg_data or lane_data). When `shape` /// spans multiple distribution units, the distributed shape may contain /// multiple such units. /// For wrap-around dimensions where the division is uneven, the tensor tile /// is broadcast to all subgroups/lanes. FailureOr> computeDistributedShape(SmallVector shape) const; /// Checks if the given shape can be distributed by the layout bool isDistributable(SmallVector shape) const; /// Check if this layout is a slice of another layout. bool isSliceOf(const xegpu::DistributeLayoutAttr& other) const; /// Check if this layout is a transpose of /// the other layout according to given permutation of (0...n-1). bool isTransposeOf(const xegpu::DistributeLayoutAttr& other, ArrayRef perm, xegpu::LayoutKind kind) const; /// Check if this layout is compatible with another layout /// at a specific level of the layout hierarchy regarding a given shape. bool isCompatibleWith(const xegpu::DistributeLayoutAttr& other, SmallVector shape, xegpu::LayoutKind level) const; /// Check if this layout is equal to another layout. /// For LayoutAttr, this compares all fields. /// For SliceAttr, this requires the same parent and same sliced dims. bool isEqualTo(const xegpu::DistributeLayoutAttr& other) const; }; } // namespace mlir::xegpu namespace mlir::xegpu::detail { template struct DistributeLayoutAttrTrait : public ::mlir::AttributeInterface::Trait { }; } // namespace mlir::xegpu::detail namespace mlir::xegpu { template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isForWorkgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).isForWorkgroup(); } template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isForSubgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).isForSubgroup(); } template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isForLane(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).isForLane(); } template int64_t detail::DistributeLayoutAttrInterfaceTraits::Model::getRank(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getRank(); } template int64_t detail::DistributeLayoutAttrInterfaceTraits::Model::getNumSubgroups(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getNumSubgroups(); } template DenseI32ArrayAttr detail::DistributeLayoutAttrInterfaceTraits::Model::getOrder(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getOrder(); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::Model::getEffectiveSgLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getEffectiveSgLayoutAsInt(); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::Model::getEffectiveSgDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getEffectiveSgDataAsInt(); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::Model::getEffectiveInstDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getEffectiveInstDataAsInt(); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::Model::getEffectiveLaneLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getEffectiveLaneLayoutAsInt(); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::Model::getEffectiveLaneDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getEffectiveLaneDataAsInt(); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::Model::getEffectiveOrderAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).getEffectiveOrderAsInt(); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::dropSgLayoutAndData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).dropSgLayoutAndData(); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::dropInstData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return (::llvm::cast(tablegen_opaque_val)).dropInstData(); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::setUnitDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims) { return (::llvm::cast(tablegen_opaque_val)).setUnitDimData(unitDims); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::setUnitDimLayout(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims) { return (::llvm::cast(tablegen_opaque_val)).setUnitDimLayout(unitDims); } template FailureOr> detail::DistributeLayoutAttrInterfaceTraits::Model::delinearizeId(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId) { return (::llvm::cast(tablegen_opaque_val)).delinearizeId(builder, loc, linearId); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::setDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, int64_t sgData, int64_t instData, int64_t laneData) { return (::llvm::cast(tablegen_opaque_val)).setDimData(dim, sgData, instData, laneData); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::dropDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup) { return (::llvm::cast(tablegen_opaque_val)).dropDims(dimGroup); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::collapseDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup) { return (::llvm::cast(tablegen_opaque_val)).collapseDims(dimGroup); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::expandDim(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, ArrayRef targetShape) { return (::llvm::cast(tablegen_opaque_val)).expandDim(dim, targetShape); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::Model::transposeDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, ArrayRef permutation) { return (::llvm::cast(tablegen_opaque_val)).transposeDims(permutation); } template FailureOr>> detail::DistributeLayoutAttrInterfaceTraits::Model::computeDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId, ArrayRef shape) { return (::llvm::cast(tablegen_opaque_val)).computeDistributedCoords(builder, loc, linearId, shape); } template SmallVector> detail::DistributeLayoutAttrInterfaceTraits::Model::computeStaticDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t linearId, ArrayRef shape) { return (::llvm::cast(tablegen_opaque_val)).computeStaticDistributedCoords(linearId, shape); } template FailureOr> detail::DistributeLayoutAttrInterfaceTraits::Model::computeDistributedShape(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape) { SmallVector layout; SmallVector subShape; if ((::llvm::cast(tablegen_opaque_val)).isForWorkgroup()) { layout = (::llvm::cast(tablegen_opaque_val)).getEffectiveSgLayoutAsInt(); subShape = (::llvm::cast(tablegen_opaque_val)).getEffectiveSgDataAsInt(); } else if ((::llvm::cast(tablegen_opaque_val)).isForSubgroup()) { layout = (::llvm::cast(tablegen_opaque_val)).getEffectiveLaneLayoutAsInt(); subShape = (::llvm::cast(tablegen_opaque_val)).getEffectiveLaneDataAsInt(); } else { return failure(); } // sgdata or lanedata cannot be empty for distributed shape computation if (subShape.empty()) return failure(); SmallVector distributedShape(shape.size()); for (auto [i, dim] : llvm::enumerate(shape)) { int64_t distriUnit = layout[i]*subShape[i]; if ((dim % distriUnit) == 0) { // Evenly divisible case, divide the dimension by the layout factor. distributedShape[i] = dim / layout[i]; if (distributedShape[i] % subShape[i] != 0) return failure(); } else { // wrap around case, the dimension size must be equal to subShape value if(dim != subShape[i]) return failure(); distributedShape[i] = dim; } } return distributedShape; } template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isDistributable(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape) { DistributeLayoutAttr curLayoutAttr = (::llvm::cast(tablegen_opaque_val)); SmallVector curShape = shape; // Phase 1: Distribute across subgroups (sg_layout + sg_data). if (curLayoutAttr.isForWorkgroup()) { auto maybeSgShape = curLayoutAttr.computeDistributedShape(curShape); if (failed(maybeSgShape)) return false; curShape = maybeSgShape.value(); curLayoutAttr = curLayoutAttr.dropSgLayoutAndData(); if (!curLayoutAttr) return true; } // Phase 2: Distribute across instruction data (inst_data). if (curLayoutAttr.isForSubgroup() && !curLayoutAttr.isForLane()) { SmallVector instData = curLayoutAttr.getEffectiveInstDataAsInt(); for (size_t i = 0; i < curShape.size(); ++i) { if (curShape[i] % instData[i] != 0) return false; } // inst_data becomes the new shape for next phase curShape = instData; curLayoutAttr = curLayoutAttr.dropInstData(); if (!curLayoutAttr) return true; } // Phase 3: Distribute across lanes (lane_layout + lane_data). auto maybeLaneShape = curLayoutAttr.computeDistributedShape(curShape); return succeeded(maybeLaneShape); } template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isSliceOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other) { return (::llvm::cast(tablegen_opaque_val)).isSliceOf(other); } template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isTransposeOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, ArrayRef perm, xegpu::LayoutKind kind) { return (::llvm::cast(tablegen_opaque_val)).isTransposeOf(other, perm, kind); } template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isCompatibleWith(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, SmallVector shape, xegpu::LayoutKind level) { return (::llvm::cast(tablegen_opaque_val)).isCompatibleWith(other, shape, level); } template bool detail::DistributeLayoutAttrInterfaceTraits::Model::isEqualTo(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other) { return (::llvm::cast(tablegen_opaque_val)).isEqualTo(other); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isForWorkgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->isForWorkgroup(tablegen_opaque_val); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isForSubgroup(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->isForSubgroup(tablegen_opaque_val); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isForLane(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->isForLane(tablegen_opaque_val); } template int64_t detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getRank(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getRank(tablegen_opaque_val); } template int64_t detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getNumSubgroups(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getNumSubgroups(tablegen_opaque_val); } template DenseI32ArrayAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getOrder(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getOrder(tablegen_opaque_val); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getEffectiveSgLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getEffectiveSgLayoutAsInt(tablegen_opaque_val); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getEffectiveSgDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getEffectiveSgDataAsInt(tablegen_opaque_val); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getEffectiveInstDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getEffectiveInstDataAsInt(tablegen_opaque_val); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getEffectiveLaneLayoutAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getEffectiveLaneLayoutAsInt(tablegen_opaque_val); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getEffectiveLaneDataAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getEffectiveLaneDataAsInt(tablegen_opaque_val); } template SmallVector detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::getEffectiveOrderAsInt(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->getEffectiveOrderAsInt(tablegen_opaque_val); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::dropSgLayoutAndData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->dropSgLayoutAndData(tablegen_opaque_val); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::dropInstData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val) { return static_cast(impl)->dropInstData(tablegen_opaque_val); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::setUnitDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims) { return static_cast(impl)->setUnitDimData(tablegen_opaque_val, unitDims); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::setUnitDimLayout(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const SmallVector unitDims) { return static_cast(impl)->setUnitDimLayout(tablegen_opaque_val, unitDims); } template FailureOr> detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::delinearizeId(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId) { return static_cast(impl)->delinearizeId(tablegen_opaque_val, builder, loc, linearId); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::setDimData(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, int64_t sgData, int64_t instData, int64_t laneData) { return static_cast(impl)->setDimData(tablegen_opaque_val, dim, sgData, instData, laneData); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::dropDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup) { return static_cast(impl)->dropDims(tablegen_opaque_val, dimGroup); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::collapseDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector dimGroup) { return static_cast(impl)->collapseDims(tablegen_opaque_val, dimGroup); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::expandDim(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t dim, ArrayRef targetShape) { return static_cast(impl)->expandDim(tablegen_opaque_val, dim, targetShape); } template xegpu::DistributeLayoutAttr detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::transposeDims(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, ArrayRef permutation) { return static_cast(impl)->transposeDims(tablegen_opaque_val, permutation); } template FailureOr>> detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::computeDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, OpBuilder & builder, Location loc, Value linearId, ArrayRef shape) { return static_cast(impl)->computeDistributedCoords(tablegen_opaque_val, builder, loc, linearId, shape); } template SmallVector> detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::computeStaticDistributedCoords(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, int64_t linearId, ArrayRef shape) { return static_cast(impl)->computeStaticDistributedCoords(tablegen_opaque_val, linearId, shape); } template FailureOr> detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::computeDistributedShape(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape) { return static_cast(impl)->computeDistributedShape(tablegen_opaque_val, shape); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isDistributable(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, SmallVector shape) { return static_cast(impl)->isDistributable(tablegen_opaque_val, shape); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isSliceOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other) { return static_cast(impl)->isSliceOf(tablegen_opaque_val, other); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isTransposeOf(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, ArrayRef perm, xegpu::LayoutKind kind) { return static_cast(impl)->isTransposeOf(tablegen_opaque_val, other, perm, kind); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isCompatibleWith(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other, SmallVector shape, xegpu::LayoutKind level) { return static_cast(impl)->isCompatibleWith(tablegen_opaque_val, other, shape, level); } template bool detail::DistributeLayoutAttrInterfaceTraits::FallbackModel::isEqualTo(const Concept *impl, ::mlir::Attribute tablegen_opaque_val, const xegpu::DistributeLayoutAttr& other) { return static_cast(impl)->isEqualTo(tablegen_opaque_val, other); } } // namespace mlir::xegpu