/*===- TableGen'erated file -------------------------------------*- C++ -*-===*\ |* *| |* Op Declarations *| |* *| |* Automatically generated file, do not edit! *| |* From: XeGPU.td *| |* *| \*===----------------------------------------------------------------------===*/ namespace mlir::xegpu { /// It allocates a set of named barriers. /// AllocNbarrier is to create a set of named barriers as /// specified by `nbarrier_num`. Named barriers are workgroup level resources, /// and are shared by all threads in the workgroup. For example, there are /// up to 32 barriers (range 0-31) for each XeCore on PVC. A typical use case /// is that a workgroup is partitioned into N subgroups of threads (N <= 32), /// and each subgroup coordinating their work with a separate barrier with id /// range from 0 to N respectively. class AllocNbarrierOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// Atomic read-modify-write operation on the TensorDesc. /// The `xegpu.atomic_rmw` operation provides a way to perform a read-modify-write /// operation on the region described by the `TensorDesc` free from data races. The /// `kind` enumeration specifies the modification to be performed, The `mask` operand /// has the same shape with `TensorDesc`, and is used to enable or disable specific /// data points of the `TensorDesc`. The `value` operand represents the new value to /// be applied during the modification. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// - `kind`: An attribute that specifies the atomic operation to be performed /// (e.g., add, min, max, exchange, etc.). /// /// - `tensorDesc`: A `TensorDesc` describing the memory region on which the atomic /// read-modify-write is performed. /// /// - `mask`: A predicate mask with the same shape as `tensorDesc`. Only elements /// with a true (non-zero) mask value participate in the atomic operation; /// masked-out elements are not modified. /// /// - `value`: The input values used by the atomic operation. It must have the same /// shape and element type as `tensorDesc` and `result`. /// /// - `layout`: [optional] An attribute that identifies the operation as an anchor, /// enabling users to assign a layout that governs distribution at the subgroup /// and/or lane level. Only valid at workgroup and subgroup levels. class AtomicRMWOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// Convert the layout of the input operand /// `convert_layout` redistribute data across subgroups and/or lanes from the `input_layout` to /// the `target_layout`. Both `input_layout` and `target_layout` must correspond to the same programming /// scope, such as workgroup level (wg) or subgroup level (sg) code. This operation is not valid once /// the IR is lowered to WI level because that is the end result of all distributions. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// - `source`: The input vector whose data is to be redistributed. The source and /// result types must match. /// - `input_layout`: The layout attribute describing the current distribution of `source` /// across subgroups and/or lanes. /// - `target_layout`: The layout attribute describing the desired distribution of the result /// across subgroups and/or lanes. /// /// Example (Subgroup level): /// ```mlir /// %coop_a = xegpu.convert_layout %a <{ /// input_layout = #xegpu.layout, /// target_layout = #xegpu.layout}> /// : vector<128x128xf16> /// ``` class ConvertLayoutOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// Create a memory descriptor. /// Creates a memory descriptor from a shared local memory (SLM) buffer, and xegpu /// specific memory layout. The resulting memory descriptor has to have the same size /// as the underlying shared local memory. /// /// Arguments: /// - `source` : 1D or 2D statically shaped memref, representing the raw SLM buffer. The provided memref must be contiguous. /// /// Results: /// - `mem_desc` : the memory descriptor (1D or higher). /// /// Example (1D): /// ```mlir /// %mdesc = xegpu.create_mem_desc %mref /// : memref<128xi8, 3> -> !xegpu.mem_desc<64xf16> /// ``` /// /// Example (2D with layout): /// ```mlir /// %mdesc = xegpu.create_mem_desc %mref /// : memref<4096xi8, 3> /// -> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout> /// ``` class CreateMemDescOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// Create nd-tensor descriptor operation /// The "create_nd_tdesc" operation creates a TensorDescType which represents /// a sub-view of a 1D/2D memory region inside the one or two innermost dimensions /// of the source. (It can be extended to support n-D memory region if needed in /// future). Elements in the subview continuous in each dimension. It encodes the /// following important information for supporting Intel hardware features: /// /// Arguments: /// - `source`: an object representing (starting address/pointer of) a memory region. /// It can be either a memref object, or simply a pointer represented by uint64_t type. /// For the case of dynamic memrefs or pointer, the shape and layout information of the /// memory region should be explicitly passed via `shape` and `strides` parameters. /// /// - `shape`: the shape information of the memory region pointed by the "source". It is /// typically encoded via the MemRefType of the source, e.g., memref<4096x4096xf16>. /// But if "source" is simply a pointer represented as uint64_t type, or a memref /// type without shape information e.g., memref, the shape information has /// to be explicitly passed via the "shape" and "const_shape" arguments. /// /// - `strides`: the strides of the memory region pointed by the "source". Similar to shape, /// it is typically encoded via the MemRefType of the source too. But if "source" is /// simply a pointer represented as uint64_t type, or a memref type without shape /// information e.g., memref, the strides information has to be explicitly /// passed via the "strides" and "const_strides" argument. /// /// Results: /// - `res`: nd tensor descriptor /// /// Example 1 (suppose the tensor shape inferred by the compiler is 8x16): /// ```mlir /// %0 = memref.alloc() : memref<1024x1024xf32> /// %1 = xegpu.create_nd_tdesc %0 : memref<1024x1024xf32> -> TensorDesc<8x16xf32> /// ``` /// /// Example 2 (suppose the tensor shape inferred by the compiler is 8x16): /// ```mlir /// %0 = memref.alloc(%h, %w) : memref /// %c1 = arith.constant 1 : index /// %1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: memref -> TensorDesc<8x16xf32> /// ``` /// /// Example 3 (suppose the tensor shape inferred by the compiler is 8x16): /// ```mlir /// %0 = ... : ui64 /// %c1 = arith.constant 1 : index /// %1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: ui64 -> TensorDesc<8x16xf32> /// ``` class CreateNdDescOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// It performs scaled mma computation /// DPAS MX performs matrix multiplication on matrix A and Matrix B /// of low precision data type. A is of size`mxk` /// size, B is of size `kxn`, and accumulate on matrix C of size `mxn` to the same size /// matrix. /// /// The operands can be 2D, 3D, or 4D vectors. When the vectors have more than 2 /// dimensions, the leading dimensions are treated as **batch dimensions** that /// must match across all operands (a, b, acc, result, scale_a, scale_b). The batch /// dimensions represent independent matrix multiplications that are executed in /// parallel. For example: /// - 2D: `A: vector<8x32xf8E5M2>`, `B: vector<32x16xf8E5M2>` -> `C: vector<8x16xbf16>` /// - 3D: `A: vector<4x8x32xf8E5M2>`, `B: vector<4x32x16xf8E5M2>` -> `C: vector<4x8x16xbf16>` /// (4 independent 8x16 scaled matrix multiplications) /// - 4D: `A: vector<2x4x8x32xf8E5M2>`, `B: vector<2x4x32x16xf8E5M2>` -> `C: vector<2x4x8x16xbf16>` /// (2x4=8 independent 8x16 scaled matrix multiplications) /// /// The last 2 dimensions are always the core matrix multiplication dimensions /// (M, K for a; K, N for b; M, N for result). The scale vectors also follow the /// same batch dimension structure. /// /// In lane level code, each lane from a subgroup holds a data fragment for A, B, Acc and the result, /// which are represented as 1D vectors. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// /// - `a`: A vector value representing the left-hand-side matrix tile (A) participating in the /// matrix multiply. Can be 1D, 2D, 3D, or 4D where leading dimensions are batch dimensions. /// /// - `b`: A vector value representing the right-hand-side matrix tile (B). Can be 1D, 2D, 3D, or 4D /// where leading dimensions are batch dimensions. /// /// - `acc`: A vector value representing the accumulator matrix tile (C). The /// result is computed as `a * b + acc`. Must have the same batch dimensions as a and b. /// /// - `scale_a`: A floating point vector/scalar value used to scale `a` for /// matrix multiplication. When a vector, must have matching batch dimensions. /// /// - `scale_b`: A floating point vector/scalar value used to scale `b` for /// matrix multiplication. When a vector, must have matching batch dimensions. /// /// - `layout_a`, `layout_b`, `layout_cd`: [optional] Attributes that identify this /// operation as anchor for operands A, B, and the accumulator/result, enabling users to assign layouts /// that govern distribution at the subgroup and/or lane level. Only valid at workgroup and subgroup /// level. class DpasMxOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// It performs mma computation /// DPAS performs matrix multiplication on matrix A of `mxk` /// size, B of `kxn` size, and accumulate on matrix C of `mxn` to the same size /// matrix. /// /// The operands can be 2D, 3D, or 4D vectors. When the vectors have more than 2 /// dimensions, the leading dimensions are treated as **batch dimensions** that /// must match across all operands (lhs, rhs, acc, and result). The batch /// dimensions represent independent matrix multiplications that are executed in /// parallel. For example: /// - 2D: `A: vector<8x16xf16>`, `B: vector<16x16xf16>` -> `C: vector<8x16xf32>` /// - 3D: `A: vector<4x8x16xf16>`, `B: vector<4x16x16xf16>` -> `C: vector<4x8x16xf32>` /// (4 independent 8x16 matrix multiplications) /// - 4D: `A: vector<2x4x8x16xf16>`, `B: vector<2x4x16x16xf16>` -> `C: vector<2x4x8x16xf32>` /// (2x4=8 independent 8x16 matrix multiplications) /// /// The last 2 dimensions are always the core matrix multiplication dimensions /// (M, K for lhs; K, N for rhs; M, N for result). /// /// In lane level code, each lane from a subgroup holds a data fragment for A, B, C and the result, /// which are represented as 1D vectors. Please refer to [OpenCL Intel extentions] /// (https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html) /// for more details about the fragment distribution. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// /// - `lhs`: A vector value representing the left-hand-side matrix tile (A) participating in the /// matrix multiply. Can be 1D, 2D, 3D, or 4D where leading dimensions are batch dimensions. /// /// - `rhs`: A vector value representing the right-hand-side matrix tile (B). Can be 1D, 2D, 3D, or 4D /// where leading dimensions are batch dimensions. /// /// - `acc`: [optional] A vector value representing the accumulator matrix tile (C). When present, the /// result is computed as `lhs * rhs + acc`; otherwise, the accumulator is implicitly assumed to be zero. /// Must have the same batch dimensions as lhs and rhs. /// /// - `layout_a`, `layout_b`, `layout_cd`: [optional] Attributes that identify this /// operation as anchor for operands A, B, and the accumulator/result, enabling users to assign layouts /// that govern distribution at the subgroup and/or lane level. Only valid at workgroup and subgroup /// level. /// /// Example 1 (Workgroup level, 2D): /// /// ```mlir /// %d = xegpu.dpas %a, %b, %c <{ /// layout_a = #xegpu.layout, /// layout_b = #xegpu.layout, /// layout_cd = #xegpu.layout} /// : vector<64x128xf16>, vector<128x128xf16>, vector<64x128xf32> -> vector<64x128xf32> /// ``` /// /// Example 2 (Lane level, 1D): /// /// ```mlir /// %d = xegpu.dpas %a, %b, %c /// : vector<8xf16>, vector<16xf16>, vector<8xf32> -> vector<8xf32> /// ``` /// /// Example 3 (Workgroup level, 3D with batch): /// /// ```mlir /// // 4 independent 8x16 matrix multiplications /// %d = xegpu.dpas %a, %b, %c /// : vector<4x8x16xf16>, vector<4x16x16xf16>, vector<4x8x16xf32> -> vector<4x8x16xf32> /// ``` class DpasOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// It synchronizes memory accesses. /// It synchronizes the memory access between /// write and following read or write. /// 1. `Memory_kind` describes the memory kind. "global" means the global memory, /// "slm" means the share local memory. /// 2. `Fence_scope` describes the scope of fence. "Workgroup" means that the scope would be /// within each workgroup. "GPU" means the scope would be across workgroups within the GPU. class FenceOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// It assigns a named barrier to the current thread. /// InitNbarrierOp assigns the named barrier with the specified /// barrier ID (0~31) to the current thread. Multiple threads may bind to the /// same named barrier, and the `participant_thread_num` specifies the total /// number of threads associated with the nbarrier. It returns an object of /// NbarrierType representing the barrier class InitNbarrierOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// load a set of scattered data points from memory. /// It (aka. load) load data per each lane. The output /// describes the data being loaded at the subgroup level, so its size is /// consistent with the number of lanes in a subgroup. When the chunk size /// is larger than 2, the output vector is a 2D vector, with dim-0 correspoding /// to lanes, and dim-1 corresponding to the chunk size loaded by each lane. /// The mask operand masks out memory access so that it is safe to pass out-of-boundary /// addresses/offsets as long as they are masked. Each mask element applies to one lane. /// /// In lane level, the result is a 1D vector that represents the data to be loaded by /// each lane. If size is not 1, size should be equal to the chunk size. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// /// - `source`: represents the memory region to be loaded from, which can be either a /// 1D memref or pointer (ui64, ui32, i64 or i32). /// /// - `offsets`: represents offsets from source. /// offsets is a vector of `index` type and vector length is either the subgroup size /// or 1 at lane level. scalar offset is also valid for lane level. /// /// - `mask`: is a vector of `i1` type, which is used to mask out the memory access. /// mask is a vector of size equal to the subgroup size, or 1 at lane level. /// scalar mask is also valid for lane level. /// /// - `chunk_size`: [optional] represents contiguous number of elements to load from per work item. /// /// - `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache. /// /// - `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the result /// of load. Only valid at workgroup and subgroup levels. /// /// Results: /// - `res`: represents loaded data /// /// Example 1 (Subgroup level): /// A variant accepts memref as base pointer or the source operand /// could be a raw pointer (ui64, ui32, i64, i32). /// ```mlir /// %a = memref.alloc() : memref<1024xf32> /// %offsets = vector.step : vector<16xindex> /// %mask = vector.constant_mask [16]: vector<16xi1> /// %val = xegpu.load %a[%offsets], %mask {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint, /// layout = #xegpu.layout} /// : memref<1024xf32>, vector<16xi1>, vector<16xindex> -> vector<16xf32> /// ``` /// /// Example 2 (lane level): /// lane level only accepts the offsets variant. chunk_size can be inferred from result /// type. In this example, chunk_size is 8. /// ```mlir /// %2 = xegpu.load %1[%2], %0 <{l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint}> /// : memref<128xf32>, vector<1xindex>, vector<1xi1> -> vector<8xf32> /// ``` class LoadGatherOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// This operation loads a block of data from shared local memory (SLM) as specified /// by the provided `mem_desc`. Memory descriptors of any rank (1D or higher) are supported. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// - `mem_desc`: the memory descriptor identifying the SLM region. /// - `offsets`: the coordinates within the memory descriptor to read from. /// - `subgroup_block_io`: [optional] An attribute indicating that the operation can be lowered /// to a subgroup block load. When this attribute is present, the offsets are subgroup-uniform /// across all lanes. Only used on subgroup and lane level. /// - `layout`: [optional] Describes the expected layout of the `mem_desc` operand as well as /// the result of load (they are identical). /// Only valid at workgroup and subgroup levels. /// /// Results: /// - `res`: the elements loaded from SLM. /// /// Example (1D): /// ```mlir /// %1 = xegpu.load_matrix %0[%c0] : !xegpu.mem_desc<64xf16>, index -> vector<16xf16> /// ``` /// /// Example (Workgroup level, 2D): /// ```mlir /// %c0 = arith.constant 0 : index /// %1 = xegpu.load_matrix %0[%c0, %c0] <{ /// layout = #xegpu.layout }> /// : !xegpu.mem_desc<128x128xf16, #xegpu.mem_layout> /// , index, index -> vector<128x128xf16> /// ``` class LoadMatrixOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// loads a n-D block from memory (represented by TensorDesc)to registers (represented by vector) /// LoadNdOp essentially mimics the hardware block read instruction to read /// a block of data from memory to register. It takes a set of optional cache /// hints for each level of cache, L1, L2 and L3. If hardware does not have a /// correspoding cache, Corresponding cache hint attribute will be masked. /// /// The tensor descriptor can be 1D, 2D, 3D, or higher dimensional. When the /// tensor descriptor has more than 2 dimensions, the leading dimensions are /// treated as **batch dimensions** that are unrolled to unit dimensions during /// lowering. At the subgroup level, the load_nd operation executes at 2D /// granularity to match the 2D block IO hardware support. The result vector /// has the same shape as the tensor descriptor. The number of offset indices /// must match the rank of the tensor descriptor. /// /// On Intel GPUs, hardware-supported packing rearranges data elements during /// the load of the B operand when the element bit-width is less than 32 bits /// (for example, fp16). The transpose feature reorders data during the load /// when the element type is fp32 or fp64. These two features are mutually /// exclusive and shall not be enabled simultaneously. Both features support only /// 2D blocked tensor_desc. /// /// At lane level, result vector represents the data to be loaded by each lane. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// /// - `TensorDesc`: A tensor descriptor specifying the base nd-region of memory /// and the tensor tile to be loaded. Can be 1D, 2D, 3D, or higher dimensional /// where leading dimensions are batch dimensions. /// /// - `offsets`: Index values representing per-dimension offsets from the base position /// encoded in `TensorDesc`. They are encoded via `offsets` and `const_offsets`. /// The number of offsets must match the tensor descriptor rank. /// /// - `packed`: [optional] A unit attribute indicating that packing is applied /// during the load when supported by the hardware. Only valid at lane level. /// /// - `transpose`: [optional] An attribute describing a hardware-supported transpose /// to be applied during the load. Only valid at Lane level. /// /// - `l1_hint`, `l2_hint`, `l3_hint`: [optional] Cache-hint attributes indicating the /// desired behavior at the L1, L2, and L3 cache levels. /// /// - `layout`: [optional] Describes the expected layout of the `tensor_desc` operand as well as the result of the load (they are identical). Only valid at workgroup and subgroup levels. /// /// Example 1 (Workgroup level, 2D): /// ```mlir /// xegpu.load_nd %1 {transpose = [1, 0], /// l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint, /// layout = #xegpu.layout} /// : !xegpu.tensor_desc<32x256xf32> -> vector<32x256xf32> /// ``` /// /// Example 2 (lane level, 2D): /// ```mlir /// xegpu.load_nd %1 {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint}> /// : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32> /// ``` /// /// Example 3 (3D with batch dimension): /// ```mlir /// // Load 4 independent 8x16 blocks /// %result = xegpu.load_nd %tdesc[0, 0, 0] <{l1_hint = #xegpu.cache_hint}> /// : !xegpu.tensor_desc<4x8x16xf16> -> vector<4x8x16xf16> /// ``` class LoadNdOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// It signals the arrival at the named barrier. /// NbarrierArriveOp signals the hardware (or other threads) /// that the current thread has produced its data for the consumer threads. When /// the hardware signalled by `participant_thread_num` threads for the named barrier, /// it will notify the threads waiting for the named barrier to continue their work. class NbarrierArriveOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// It waits for a named barrier. /// NbarrierWaitOp signals the hardware which named barrier /// the current thread is waiting for, such that it can get notified when the /// named barrier is completed. class NbarrierWaitOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// prefetches a n-D block to cache /// It issues an instruction to prefetch a block of data from continuous /// memory regions to each level of the cache based on their cache policy. /// /// The tensor descriptor can be n-dimensional. When the /// tensor descriptor has more than 2 dimensions, the leading dimensions are /// treated as **batch dimensions** that are unrolled to unit dimensions during /// lowering. At the subgroup level, the prefetch_nd operation executes at 2D /// granularity to match the 2D block IO hardware support. The number of offset /// indices must match the rank of the tensor descriptor. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// - `TensorDesc`: A tensor descriptor specifying the base nd-region of /// memory and tensor tile to be prefetched. Can be 1D, 2D, 3D, or higher /// dimensional where leading dimensions are batch dimensions. /// /// - `offsets`: index values representing per-dimension offsets from the /// base position encoded in `TensorDesc`. It is encoded via "offsets" /// and "const_offsets". The number of offsets must match the tensor descriptor rank. /// /// - `l1_hint`, `l2_hint`, `l3_hint`: [optional] An cache-hint attribute /// indicating the desired behavior at the L1, L2, and L3 cache levels. /// /// - `layout`: [optional] Describes the expected layout of the `tensor_desc` operand. /// Only valid at the workgroup and subgroup levels. /// /// Example 1 (Workgroup level, 2D): /// ```mlir /// %c0 = arith.constant 0 : index /// %c1 = arith.constant 1 : index /// xegpu.prefetch_nd %tdesc[%c0, %c1] {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint, /// layout = #xegpu.layout } /// : !xegpu.tensor_desc<32x256xf16> /// ``` /// /// Example 2 (3D with batch dimension): /// ```mlir /// // Prefetch 4 independent 8x16 blocks /// xegpu.prefetch_nd %tdesc[0, 0, 0] <{l1_hint = #xegpu.cache_hint}> /// : !xegpu.tensor_desc<4x8x16xf16> /// ``` class PrefetchNdOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// prefetches a set of scattered data points to cache /// It issues instructions to prefetch a set of scattered data points /// from memory to each level of the cache based on their cache policy. /// As compared to prefetch_nd, which works on non-scattered TensorDesc, /// it works on scattered TensorDesc instead. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// /// - `source`: represents the memory region to be loaded from, which can be either a /// 1D memref or pointer (ui64, ui32, i64 or i32). /// /// - `offsets`: represents offsets from source. /// offsets is a vector of `index` type and vector length is either the subgroup size /// or 1 at lane level. scalar offset is also valid for lane level. /// /// - `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache. /// /// - `offset_align_byte`: [optional] required if `source` is a pointer. If `source` is not a pointer, /// it is not allowed. Represents the alignment in bytes of each offset in offsets. /// /// - `layout`: [optional] Describes the expected layout of the `tensor_desc` or `offsets` /// operand. Only valid at workgroup and subgroup levels. /// /// Example 1 (Workgroup level): /// ```mlir /// %a = memref.alloc() : memref<256xf16> /// %offsets = arith.constant dense<[0, 1, ..., 255]> : vector<256xindex> /// xegpu.prefetch %a[%offsets] {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint, /// layout = #xegpu.layout /// } /// : memref<256xf16>, vector<256xindex> /// ``` /// /// Example 2 (lane level): /// A variant accepts memref or integer (raw pointer) as base and offsets directly. /// The source operand could be a raw pointer (ui64, ui32, i64, i32). /// ```mlir /// %a = memref.alloc() : memref<1024xf32> /// %0 = arith.constant dense<[0, 16, 32, 64]> : vector<4xindex> /// xegpu.prefetch %a[%0] {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint} /// : memref<1024xf32>, vector<4xindex> /// ``` /// /// Example 3 (lane level): /// lane level only accepts the offsets variant. /// ```mlir /// xegpu.prefetch %0[%1] {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint} /// : memref<256xf32>, vector<1xindex> /// ``` /// /// Example 4 (lane level): /// lane level only accepts the offsets variant. /// ```mlir /// xegpu.prefetch %0[%1] {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint, /// offset_align_byte = 2} /// : i64, vector<1xindex> /// ``` class PrefetchOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// This operation stores a `data` fragment into the shared local memory region /// specified by a `mem_desc`. Memory descriptors of any rank (1D or higher) are supported. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// - `mem_desc`: the memory descriptor specifying the SLM region. /// - `offsets`: the coordinates within the memory descriptor where the data will be written. /// - `data`: the values to be stored. /// - `subgroup_block_io`: [optional] An attribute indicating that the operation can be lowered /// to a subgroup block store. When this attribute is present, the offsets are subgroup-uniform /// across all lanes. Only used on subgroup and lane level. /// - `layout`: [optional] Describes the expected layout of the `mem_desc` operand as well as /// the value to be stored (they are identical). Only valid at workgroup and subgroup levels. /// /// Example (1D): /// ```mlir /// xegpu.store_matrix %1, %0[%c0] : vector<16xf16>, !xegpu.mem_desc<64xf16>, index /// ``` /// /// Example (Workgroup level, 2D): /// ```mlir /// %c0 = arith.constant 0 : index /// xegpu.store_matrix %1, %0[%c0, %c0] <{ /// layout = #xegpu.layout }> /// : vector<128x128xf16>, !xegpu.mem_desc<128x128xf16>>, index, index /// ``` class StoreMatrixOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// stores a n-D block register region back to memory /// StoreNdOp essentially mimics the hardware block write instruction io /// write a block of data from register into the memory region as described /// by the TensorDesc. It takes a set of optional cache hints for each level /// of cache, L1, L2 and L3. If hardware does not have a correspoding cache, /// Corresponding cache hint attribute will be masked. /// /// The tensor descriptor can be 1D, 2D, 3D, or higher dimensional. When the /// tensor descriptor has more than 2 dimensions, the leading dimensions are /// treated as **batch dimensions** that are unrolled to unit dimensions during /// lowering. At the subgroup level, the store_nd operation executes at 2D /// granularity to match the 2D block IO hardware support. The value vector /// must have the same shape as the tensor descriptor. The number of offset /// indices must match the rank of the tensor descriptor. /// /// At lane level, the input vector represents the data to be stored by each lane. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// /// - `value`: A vector value representing the tensor tile to be stored. Can be /// 1D, 2D, 3D, or higher dimensional where leading dimensions are batch dimensions. /// /// - `TensorDesc`: A tensor descriptor specifying the base nd-region of memory and /// the tensor tile to be stored. Can be 1D, 2D, 3D, or higher dimensional /// where leading dimensions are batch dimensions. /// /// - `offsets`: Index values representing per-dimension offsets from the base position /// encoded in `TensorDesc`. They are encoded via `offsets` and `const_offsets`. /// The number of offsets must match the tensor descriptor rank. /// /// - `l1_hint`, `l2_hint`, `l3_hint`: [optional] Cache-hint attributes indicating the /// desired behavior at the L1, L2, and L3 cache levels. /// /// - `layout`: [optional] Describes the expected layout of the `tensor_desc` operand as well as /// the value to be stored (they are identical). Only valid at workgroup and subgroup levels. /// /// Example 1 (Workgroup level, 2D): /// ```mlir /// xegpu.store_nd %3, %2 {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint, /// layout = #xegpu.layout} /// : vector<32x256xf16>, !xegpu.tensor_desc<32x256xf16> /// ``` /// /// Example 2 (lane level, 2D): /// ```mlir /// xegpu.store_nd %3, %2 {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint} /// : vector<8xf16>, !xegpu.tensor_desc<8x16xf16> /// ``` /// /// Example 3 (3D with batch dimension): /// ```mlir /// // Store 4 independent 8x16 blocks /// xegpu.store_nd %value, %tdesc[0, 0, 0] <{l1_hint = #xegpu.cache_hint}> /// : vector<4x8x16xf16>, !xegpu.tensor_desc<4x8x16xf16> /// ``` class StoreNdOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// store data to scattered memory locations. /// It (aka. store) stores data to scattered memory locations. The value is /// typically a 1D vector. But when the chunk size of the TensorDesc is larger than 1, it will be /// a 2D vector instead. For the later case, dim-1 of the value correspods to the simd lanes /// and the dim-0 of the value corresponds to the chunk size stored per lane. So `store_scatter` /// has transpose effect, which is similar to `load_gather`. Therefore, a transpose attribute is /// introduced on purpose, making sure users are aware of this implicit transformation. /// /// In lane level, the result is a 1D vector that represents the data to be stored by /// each lane. If size is not 1, size should be equal to the chunk size. /// /// This operation serves as an anchor through which users assign a layout attribute /// to govern computation distribution. /// /// Arguments: /// /// - `value`: represents the data to be stored. /// /// - `dest`: represents the memory region to be stored to, which can be either a /// 1D memref or pointer (ui64, ui32, i64 or i32). /// /// - `offsets`: represents offsets from dest. /// offsets is a vector of `index` type and vector length is either the subgroup size /// or 1 at lane level. scalar offset is also valid for lane level. /// /// - `mask`: is a vector of `i1` type, which is used to mask out the memory access. /// mask is a vector of size equal to the subgroup size, or 1 at lane level. /// scalar mask is also valid for lane level. /// /// - `chunk_size`: [optional] represents contiguous number of elements to store to per work item. /// /// - `l1_hint`, `l2_hint`, `l3_hint`: [optional] cache hints for each level of cache. /// /// - `layout`: [optional] Describes the expected layout of the `tensor_desc` operand or the value /// to be stored. Only valid at workgroup and subgroup levels. /// /// /// Example 1 (Subgroup level): /// A variant accepts memref as base pointer and an offset. /// The dest operand could be a raw pointer (uint64_t). /// ```mlir /// %a = memref.alloc() : memref<1024xf32> /// %val = arith.constant dense<0.0> : vector<16xf32> /// %offsets = vector.step : vector<16xindex> /// %mask = vector.constant_mask [16]: vector<16xi1> /// xegpu.store %val, %a[%offsets], %mask {l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint, /// layout = #xegpu.layout} /// : vector<16xf32>, memref<1024xf32>, vector<16xi1>, vector<16xindex> /// ``` /// /// Example 2 (Lane level): /// Lane level IR only accepts the offsets variant. chunk_size can be inferred from value /// type. In this example, chunk_size is 8. /// ```mlir /// xegpu.store %0, %1[%2], %3 <{l1_hint = #xegpu.cache_hint, /// l2_hint = #xegpu.cache_hint, /// l3_hint = #xegpu.cache_hint}> /// : vector<8xf32>, memref<256xf32>, vector<1xindex>, vector<1xi1> /// ``` class StoreScatterOp; } // namespace mlir::xegpu namespace mlir::xegpu { /// It performs floating point truncation from higher precision to lower precision. /// The `xegpu.truncf` operation truncates floating point values from a higher /// precision type to a lower precision type. /// Converts `f16` and `bf16` to microscaling float types. /// Rounding mode defaults to round to nearest even. /// /// Example: /// ```mlir /// %res = xegpu.truncf %src : vector<16xf16> -> vector<16xf8E5M2> /// ``` class TruncfOp; } // namespace mlir::xegpu #ifdef GET_OP_CLASSES #undef GET_OP_CLASSES namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::AllocNbarrierOp declarations //===----------------------------------------------------------------------===// namespace detail { class AllocNbarrierOpGenericAdaptorBase { public: struct Properties { using nbarrier_numTy = ::mlir::IntegerAttr; nbarrier_numTy nbarrier_num; auto getNbarrierNum() const { auto &propStorage = this->nbarrier_num; return ::llvm::cast<::mlir::IntegerAttr>(propStorage); } void setNbarrierNum(const ::mlir::IntegerAttr &propValue) { this->nbarrier_num = propValue; } bool operator==(const Properties &rhs) const { return rhs.nbarrier_num == this->nbarrier_num && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: AllocNbarrierOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.alloc_nbarrier", odsAttrs.getContext()); } AllocNbarrierOpGenericAdaptorBase(AllocNbarrierOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::IntegerAttr getNbarrierNumAttr() { auto attr = ::llvm::cast<::mlir::IntegerAttr>(getProperties().nbarrier_num); return attr; } uint64_t getNbarrierNum(); }; } // namespace detail template class AllocNbarrierOpGenericAdaptor : public detail::AllocNbarrierOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::AllocNbarrierOpGenericAdaptorBase; public: AllocNbarrierOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} AllocNbarrierOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : AllocNbarrierOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} AllocNbarrierOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : AllocNbarrierOpGenericAdaptor(values, attrs, Properties{}, {}) {} AllocNbarrierOpGenericAdaptor(RangeT values, const AllocNbarrierOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> AllocNbarrierOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class AllocNbarrierOpAdaptor : public AllocNbarrierOpGenericAdaptor<::mlir::ValueRange> { public: using AllocNbarrierOpGenericAdaptor::AllocNbarrierOpGenericAdaptor; AllocNbarrierOpAdaptor(AllocNbarrierOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class AllocNbarrierOp : public ::mlir::Op { public: using Op::Op; using Op::print; using Adaptor = AllocNbarrierOpAdaptor; template using GenericAdaptor = AllocNbarrierOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("nbarrier_num")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getNbarrierNumAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getNbarrierNumAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.alloc_nbarrier"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::IntegerAttr getNbarrierNumAttr() { return ::llvm::cast<::mlir::IntegerAttr>(getProperties().nbarrier_num); } uint64_t getNbarrierNum(); void setNbarrierNumAttr(::mlir::IntegerAttr attr) { getProperties().nbarrier_num = attr; } void setNbarrierNum(uint64_t attrValue); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::IntegerAttr nbarrier_num); static AllocNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::IntegerAttr nbarrier_num); static AllocNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::IntegerAttr nbarrier_num); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::IntegerAttr nbarrier_num); static AllocNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::IntegerAttr nbarrier_num); static AllocNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::IntegerAttr nbarrier_num); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, uint64_t nbarrier_num); static AllocNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, uint64_t nbarrier_num); static AllocNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, uint64_t nbarrier_num); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, uint64_t nbarrier_num); static AllocNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, uint64_t nbarrier_num); static AllocNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, uint64_t nbarrier_num); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static AllocNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static AllocNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static AllocNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static AllocNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 1 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::AllocNbarrierOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::AllocNbarrierOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::AtomicRMWOp declarations //===----------------------------------------------------------------------===// namespace detail { class AtomicRMWOpGenericAdaptorBase { public: struct Properties { using kindTy = ::mlir::arith::AtomicRMWKindAttr; kindTy kind; auto getKind() const { auto &propStorage = this->kind; return ::llvm::cast<::mlir::arith::AtomicRMWKindAttr>(propStorage); } void setKind(const ::mlir::arith::AtomicRMWKindAttr &propValue) { this->kind = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } bool operator==(const Properties &rhs) const { return rhs.kind == this->kind && rhs.layout == this->layout && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: AtomicRMWOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.atomic_rmw", odsAttrs.getContext()); } AtomicRMWOpGenericAdaptorBase(AtomicRMWOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::arith::AtomicRMWKindAttr getKindAttr() { auto attr = ::llvm::cast<::mlir::arith::AtomicRMWKindAttr>(getProperties().kind); return attr; } ::mlir::arith::AtomicRMWKind getKind(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class AtomicRMWOpGenericAdaptor : public detail::AtomicRMWOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::AtomicRMWOpGenericAdaptorBase; static constexpr int odsIndex_tensorDesc = 0; static constexpr int odsIndex_mask = 1; static constexpr int odsIndex_value = 2; public: AtomicRMWOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} AtomicRMWOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : AtomicRMWOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} AtomicRMWOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : AtomicRMWOpGenericAdaptor(values, attrs, Properties{}, {}) {} AtomicRMWOpGenericAdaptor(RangeT values, const AtomicRMWOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> AtomicRMWOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getTensorDesc() { return (*getODSOperands(0).begin()); } ValueT getMask() { return (*getODSOperands(1).begin()); } ValueT getValue() { return (*getODSOperands(2).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class AtomicRMWOpAdaptor : public AtomicRMWOpGenericAdaptor<::mlir::ValueRange> { public: using AtomicRMWOpGenericAdaptor::AtomicRMWOpGenericAdaptor; AtomicRMWOpAdaptor(AtomicRMWOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class AtomicRMWOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::NOperands<3>::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::ConditionallySpeculatable::Trait, ::mlir::OpTrait::AlwaysSpeculatableImplTrait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = AtomicRMWOpAdaptor; template using GenericAdaptor = AtomicRMWOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_tensorDesc = 0; static constexpr int odsIndex_mask = 1; static constexpr int odsIndex_value = 2; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("kind"), ::llvm::StringRef("layout")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getKindAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getKindAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.atomic_rmw"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::TensorDescType> getTensorDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::TensorDescType>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::VectorType> getMask() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(1).begin()); } ::mlir::TypedValue<::mlir::VectorType> getValue() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(2).begin()); } ::mlir::OpOperand &getTensorDescMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getMaskMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getValueMutable() { auto range = getODSOperandIndexAndLength(2); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getResult() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::arith::AtomicRMWKindAttr getKindAttr() { return ::llvm::cast<::mlir::arith::AtomicRMWKindAttr>(getProperties().kind); } ::mlir::arith::AtomicRMWKind getKind(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setKindAttr(::mlir::arith::AtomicRMWKindAttr attr) { getProperties().kind = attr; } void setKind(::mlir::arith::AtomicRMWKind attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type result, ::mlir::arith::AtomicRMWKindAttr kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type result, ::mlir::arith::AtomicRMWKindAttr kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type result, ::mlir::arith::AtomicRMWKindAttr kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::arith::AtomicRMWKindAttr kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::arith::AtomicRMWKindAttr kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::arith::AtomicRMWKindAttr kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type result, ::mlir::arith::AtomicRMWKind kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type result, ::mlir::arith::AtomicRMWKind kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type result, ::mlir::arith::AtomicRMWKind kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::arith::AtomicRMWKind kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::arith::AtomicRMWKind kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static AtomicRMWOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::arith::AtomicRMWKind kind, ::mlir::Value tensorDesc, ::mlir::Value mask, ::mlir::Value value, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static AtomicRMWOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static AtomicRMWOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static AtomicRMWOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static AtomicRMWOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 2 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::AtomicRMWOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::AtomicRMWOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::ConvertLayoutOp declarations //===----------------------------------------------------------------------===// namespace detail { class ConvertLayoutOpGenericAdaptorBase { public: struct Properties { using input_layoutTy = ::mlir::xegpu::DistributeLayoutAttr; input_layoutTy input_layout; auto getInputLayout() const { auto &propStorage = this->input_layout; return ::llvm::cast<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setInputLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->input_layout = propValue; } using target_layoutTy = ::mlir::xegpu::DistributeLayoutAttr; target_layoutTy target_layout; auto getTargetLayout() const { auto &propStorage = this->target_layout; return ::llvm::cast<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setTargetLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->target_layout = propValue; } bool operator==(const Properties &rhs) const { return rhs.input_layout == this->input_layout && rhs.target_layout == this->target_layout && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: ConvertLayoutOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.convert_layout", odsAttrs.getContext()); } ConvertLayoutOpGenericAdaptorBase(ConvertLayoutOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::xegpu::DistributeLayoutAttr getInputLayoutAttr() { auto attr = ::llvm::cast<::mlir::xegpu::DistributeLayoutAttr>(getProperties().input_layout); return attr; } ::mlir::xegpu::DistributeLayoutAttr getInputLayout(); ::mlir::xegpu::DistributeLayoutAttr getTargetLayoutAttr() { auto attr = ::llvm::cast<::mlir::xegpu::DistributeLayoutAttr>(getProperties().target_layout); return attr; } ::mlir::xegpu::DistributeLayoutAttr getTargetLayout(); }; } // namespace detail template class ConvertLayoutOpGenericAdaptor : public detail::ConvertLayoutOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::ConvertLayoutOpGenericAdaptorBase; static constexpr int odsIndex_source = 0; public: ConvertLayoutOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} ConvertLayoutOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : ConvertLayoutOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} ConvertLayoutOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : ConvertLayoutOpGenericAdaptor(values, attrs, Properties{}, {}) {} ConvertLayoutOpGenericAdaptor(RangeT values, const ConvertLayoutOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> ConvertLayoutOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getSource() { return (*getODSOperands(0).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class ConvertLayoutOpAdaptor : public ConvertLayoutOpGenericAdaptor<::mlir::ValueRange> { public: using ConvertLayoutOpGenericAdaptor::ConvertLayoutOpGenericAdaptor; ConvertLayoutOpAdaptor(ConvertLayoutOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class ConvertLayoutOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::OneOperand, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = ConvertLayoutOpAdaptor; template using GenericAdaptor = ConvertLayoutOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_source = 0; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("input_layout"), ::llvm::StringRef("target_layout")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getInputLayoutAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getInputLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getTargetLayoutAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getTargetLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.convert_layout"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getSource() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(0).begin()); } ::mlir::OpOperand &getSourceMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getResult() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::xegpu::DistributeLayoutAttr getInputLayoutAttr() { return ::llvm::cast<::mlir::xegpu::DistributeLayoutAttr>(getProperties().input_layout); } ::mlir::xegpu::DistributeLayoutAttr getInputLayout(); ::mlir::xegpu::DistributeLayoutAttr getTargetLayoutAttr() { return ::llvm::cast<::mlir::xegpu::DistributeLayoutAttr>(getProperties().target_layout); } ::mlir::xegpu::DistributeLayoutAttr getTargetLayout(); void setInputLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().input_layout = attr; } void setTargetLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().target_layout = attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type result, ::mlir::Value source, ::mlir::xegpu::DistributeLayoutAttr input_layout, ::mlir::xegpu::DistributeLayoutAttr target_layout); static ConvertLayoutOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type result, ::mlir::Value source, ::mlir::xegpu::DistributeLayoutAttr input_layout, ::mlir::xegpu::DistributeLayoutAttr target_layout); static ConvertLayoutOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type result, ::mlir::Value source, ::mlir::xegpu::DistributeLayoutAttr input_layout, ::mlir::xegpu::DistributeLayoutAttr target_layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::xegpu::DistributeLayoutAttr input_layout, ::mlir::xegpu::DistributeLayoutAttr target_layout); static ConvertLayoutOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::xegpu::DistributeLayoutAttr input_layout, ::mlir::xegpu::DistributeLayoutAttr target_layout); static ConvertLayoutOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::xegpu::DistributeLayoutAttr input_layout, ::mlir::xegpu::DistributeLayoutAttr target_layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static ConvertLayoutOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static ConvertLayoutOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static ConvertLayoutOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static ConvertLayoutOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 2 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: xegpu::DistributeLayoutAttr getAnchorLayout() { return getTargetLayout(); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setTargetLayoutAttr(anchorLayout); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::ConvertLayoutOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::ConvertLayoutOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::CreateMemDescOp declarations //===----------------------------------------------------------------------===// namespace detail { class CreateMemDescOpGenericAdaptorBase { public: using Properties = ::mlir::EmptyProperties; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: CreateMemDescOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.create_mem_desc", odsAttrs.getContext()); } CreateMemDescOpGenericAdaptorBase(CreateMemDescOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } }; } // namespace detail template class CreateMemDescOpGenericAdaptor : public detail::CreateMemDescOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::CreateMemDescOpGenericAdaptorBase; static constexpr int odsIndex_source = 0; public: CreateMemDescOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} CreateMemDescOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : CreateMemDescOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} CreateMemDescOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : CreateMemDescOpGenericAdaptor(values, attrs, Properties{}, {}) {} CreateMemDescOpGenericAdaptor(RangeT values, const CreateMemDescOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> CreateMemDescOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getSource() { return (*getODSOperands(0).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class CreateMemDescOpAdaptor : public CreateMemDescOpGenericAdaptor<::mlir::ValueRange> { public: using CreateMemDescOpGenericAdaptor::CreateMemDescOpGenericAdaptor; CreateMemDescOpAdaptor(CreateMemDescOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class CreateMemDescOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::OneOperand, ::mlir::OpTrait::OpInvariants, ::mlir::ConditionallySpeculatable::Trait, ::mlir::OpTrait::AlwaysSpeculatableImplTrait, ::mlir::MemoryEffectOpInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = CreateMemDescOpAdaptor; template using GenericAdaptor = CreateMemDescOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_source = 0; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { return {}; } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.create_mem_desc"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getSource() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(0).begin()); } ::mlir::OpOperand &getSourceMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::MemDescType> getMemDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::MemDescType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type mem_desc, ::mlir::Value source); static CreateMemDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type mem_desc, ::mlir::Value source); static CreateMemDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type mem_desc, ::mlir::Value source); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value source); static CreateMemDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value source); static CreateMemDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value source); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static CreateMemDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static CreateMemDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static CreateMemDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static CreateMemDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); public: }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::CreateMemDescOp) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::CreateNdDescOp declarations //===----------------------------------------------------------------------===// namespace detail { class CreateNdDescOpGenericAdaptorBase { public: struct Properties { using const_shapeTy = ::mlir::DenseI64ArrayAttr; const_shapeTy const_shape; auto getConstShape() const { auto &propStorage = this->const_shape; return ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(propStorage); } void setConstShape(const ::mlir::DenseI64ArrayAttr &propValue) { this->const_shape = propValue; } using const_stridesTy = ::mlir::DenseI64ArrayAttr; const_stridesTy const_strides; auto getConstStrides() const { auto &propStorage = this->const_strides; return ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(propStorage); } void setConstStrides(const ::mlir::DenseI64ArrayAttr &propValue) { this->const_strides = propValue; } using operandSegmentSizesTy = std::array; operandSegmentSizesTy operandSegmentSizes; ::llvm::ArrayRef getOperandSegmentSizes() const { auto &propStorage = this->operandSegmentSizes; return propStorage; } void setOperandSegmentSizes(::llvm::ArrayRef propValue) { auto &propStorage = this->operandSegmentSizes; ::llvm::copy(propValue, propStorage.begin()); } bool operator==(const Properties &rhs) const { return rhs.const_shape == this->const_shape && rhs.const_strides == this->const_strides && rhs.operandSegmentSizes == this->operandSegmentSizes && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: CreateNdDescOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.create_nd_tdesc", odsAttrs.getContext()); } CreateNdDescOpGenericAdaptorBase(CreateNdDescOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::DenseI64ArrayAttr getConstShapeAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(getProperties().const_shape); return attr; } ::std::optional<::llvm::ArrayRef> getConstShape(); ::mlir::DenseI64ArrayAttr getConstStridesAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(getProperties().const_strides); return attr; } ::std::optional<::llvm::ArrayRef> getConstStrides(); }; } // namespace detail template class CreateNdDescOpGenericAdaptor : public detail::CreateNdDescOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::CreateNdDescOpGenericAdaptorBase; static constexpr int odsIndex_source = 0; static constexpr int odsIndex_shape = 1; static constexpr int odsIndex_strides = 2; public: CreateNdDescOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} CreateNdDescOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : CreateNdDescOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} CreateNdDescOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs) : CreateNdDescOpGenericAdaptor(values, attrs, Properties{}, {}) {} CreateNdDescOpGenericAdaptor(RangeT values, const CreateNdDescOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> CreateNdDescOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getSource() { return (*getODSOperands(0).begin()); } RangeT getShape() { return getODSOperands(1); } RangeT getStrides() { return getODSOperands(2); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class CreateNdDescOpAdaptor : public CreateNdDescOpGenericAdaptor<::mlir::ValueRange> { public: using CreateNdDescOpGenericAdaptor::CreateNdDescOpGenericAdaptor; CreateNdDescOpAdaptor(CreateNdDescOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class CreateNdDescOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::AtLeastNOperands<1>::Impl, ::mlir::OpTrait::AttrSizedOperandSegments, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::ConditionallySpeculatable::Trait, ::mlir::OpTrait::AlwaysSpeculatableImplTrait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::ViewLikeOpInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = CreateNdDescOpAdaptor; template using GenericAdaptor = CreateNdDescOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_source = 0; static constexpr int odsIndex_shape = 1; static constexpr int odsIndex_strides = 2; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("const_shape"), ::llvm::StringRef("const_strides"), ::llvm::StringRef("operandSegmentSizes")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getConstShapeAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getConstShapeAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getConstStridesAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getConstStridesAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getOperandSegmentSizesAttrName() { return (*this)->getName().getAttributeNames().back(); } static ::mlir::StringAttr getOperandSegmentSizesAttrName(::mlir::OperationName name) { return name.getAttributeNames().back(); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.create_nd_tdesc"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getSource() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(0).begin()); } ::mlir::Operation::operand_range getShape() { return getODSOperands(1); } ::mlir::Operation::operand_range getStrides() { return getODSOperands(2); } ::mlir::OpOperand &getSourceMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getShapeMutable(); ::mlir::MutableOperandRange getStridesMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::TensorDescType> getTensorDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::TensorDescType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::DenseI64ArrayAttr getConstShapeAttr() { return ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(getProperties().const_shape); } ::std::optional<::llvm::ArrayRef> getConstShape(); ::mlir::DenseI64ArrayAttr getConstStridesAttr() { return ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(getProperties().const_strides); } ::std::optional<::llvm::ArrayRef> getConstStrides(); void setConstShapeAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().const_shape = attr; } void setConstShape(::std::optional<::llvm::ArrayRef> attrValue); void setConstStridesAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().const_strides = attr; } void setConstStrides(::std::optional<::llvm::ArrayRef> attrValue); ::mlir::Attribute removeConstShapeAttr() { auto attr = getProperties().const_shape; getProperties().const_shape = {}; return attr; } ::mlir::Attribute removeConstStridesAttr() { auto attr = getProperties().const_strides; getProperties().const_strides = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Type tdesc, TypedValue source); static CreateNdDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Type tdesc, TypedValue source); static CreateNdDescOp create(::mlir::ImplicitLocOpBuilder &builder, Type tdesc, TypedValue source); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Type tdesc, Value source, llvm::ArrayRef shape, llvm::ArrayRef strides); static CreateNdDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Type tdesc, Value source, llvm::ArrayRef shape, llvm::ArrayRef strides); static CreateNdDescOp create(::mlir::ImplicitLocOpBuilder &builder, Type tdesc, Value source, llvm::ArrayRef shape, llvm::ArrayRef strides); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type TensorDesc, ::mlir::Value source, ::mlir::ValueRange shape, ::mlir::ValueRange strides, /*optional*/::mlir::DenseI64ArrayAttr const_shape, /*optional*/::mlir::DenseI64ArrayAttr const_strides); static CreateNdDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type TensorDesc, ::mlir::Value source, ::mlir::ValueRange shape, ::mlir::ValueRange strides, /*optional*/::mlir::DenseI64ArrayAttr const_shape, /*optional*/::mlir::DenseI64ArrayAttr const_strides); static CreateNdDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type TensorDesc, ::mlir::Value source, ::mlir::ValueRange shape, ::mlir::ValueRange strides, /*optional*/::mlir::DenseI64ArrayAttr const_shape, /*optional*/::mlir::DenseI64ArrayAttr const_strides); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::ValueRange shape, ::mlir::ValueRange strides, /*optional*/::mlir::DenseI64ArrayAttr const_shape, /*optional*/::mlir::DenseI64ArrayAttr const_strides); static CreateNdDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::ValueRange shape, ::mlir::ValueRange strides, /*optional*/::mlir::DenseI64ArrayAttr const_shape, /*optional*/::mlir::DenseI64ArrayAttr const_strides); static CreateNdDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::ValueRange shape, ::mlir::ValueRange strides, /*optional*/::mlir::DenseI64ArrayAttr const_shape, /*optional*/::mlir::DenseI64ArrayAttr const_strides); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static CreateNdDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static CreateNdDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static CreateNdDescOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static CreateNdDescOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 2 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } /// Returns the type of the source memref operand. Type getSourceType() { return getSource().getType(); } /// Returns the type of the result TensorDesc. xegpu::TensorDescType getType() { return getTensorDesc().getType(); } /// Return the element type of the TensorDesc Type getElementType() { return getType().getElementType(); } /// Return the shape of the TensorDesc llvm::ArrayRef getTensorDescShape() { return getType().getShape(); } SmallVector getMixedSizes() { SmallVector statics; /// Get the static sizes/shape, the value passed to const_shape /// will overide the value in memref shape. if (auto memrefTy = llvm::dyn_cast(getSourceType())) statics = llvm::to_vector(memrefTy.getShape()); if (auto attr = getConstShapeAttr()) statics = llvm::to_vector(attr.asArrayRef()); return getMixedValues(statics, getShape(), getContext()); } SmallVector getMixedStrides() { SmallVector statics; /// Get the static strides, the value passed to const_strides /// will overide the value in memref. if (auto memrefTy = llvm::dyn_cast(getSourceType())) statics = memrefTy.getStridesAndOffset().first; if (auto attr = getConstStridesAttr()) statics = llvm::to_vector(attr.asArrayRef()); return getMixedValues(statics, getStrides(), getContext()); } /// Return the number of leading operands before the `offsets`, /// `shape` and `strides` operands. static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } mlir::Value getViewSource() { return getSource(); } unsigned getSourceMemorySpace() { auto srcTy = getSourceType(); if (auto memrefTy = llvm::dyn_cast(srcTy)) { auto attr = memrefTy.getMemorySpace(); if (attr) { if (auto intAttr = llvm::dyn_cast(attr)) { return static_cast(intAttr.getInt()); } if (auto memSpaceAttr = llvm::dyn_cast(attr)) return static_cast(memSpaceAttr.getValue()); } } // take global as default memory scope. return static_cast(MemorySpace::Global); } xegpu::DistributeLayoutAttr getDescLayoutAttr() { return dyn_cast_if_present(getType().getLayout()); } ArrayRef getDataShape() { return getTensorDescShape(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::CreateNdDescOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::CreateNdDescOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::DpasMxOp declarations //===----------------------------------------------------------------------===// namespace detail { class DpasMxOpGenericAdaptorBase { public: struct Properties { using layout_aTy = ::mlir::xegpu::DistributeLayoutAttr; layout_aTy layout_a; auto getLayoutA() const { auto &propStorage = this->layout_a; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutA(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_a = propValue; } using layout_a_scaleTy = ::mlir::xegpu::DistributeLayoutAttr; layout_a_scaleTy layout_a_scale; auto getLayoutAScale() const { auto &propStorage = this->layout_a_scale; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutAScale(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_a_scale = propValue; } using layout_bTy = ::mlir::xegpu::DistributeLayoutAttr; layout_bTy layout_b; auto getLayoutB() const { auto &propStorage = this->layout_b; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutB(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_b = propValue; } using layout_b_scaleTy = ::mlir::xegpu::DistributeLayoutAttr; layout_b_scaleTy layout_b_scale; auto getLayoutBScale() const { auto &propStorage = this->layout_b_scale; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutBScale(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_b_scale = propValue; } using layout_cdTy = ::mlir::xegpu::DistributeLayoutAttr; layout_cdTy layout_cd; auto getLayoutCd() const { auto &propStorage = this->layout_cd; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutCd(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_cd = propValue; } using operandSegmentSizesTy = std::array; operandSegmentSizesTy operandSegmentSizes; ::llvm::ArrayRef getOperandSegmentSizes() const { auto &propStorage = this->operandSegmentSizes; return propStorage; } void setOperandSegmentSizes(::llvm::ArrayRef propValue) { auto &propStorage = this->operandSegmentSizes; ::llvm::copy(propValue, propStorage.begin()); } bool operator==(const Properties &rhs) const { return rhs.layout_a == this->layout_a && rhs.layout_a_scale == this->layout_a_scale && rhs.layout_b == this->layout_b && rhs.layout_b_scale == this->layout_b_scale && rhs.layout_cd == this->layout_cd && rhs.operandSegmentSizes == this->operandSegmentSizes && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: DpasMxOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.dpas_mx", odsAttrs.getContext()); } DpasMxOpGenericAdaptorBase(DpasMxOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::xegpu::DistributeLayoutAttr getLayoutAAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_a); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutA(); ::mlir::xegpu::DistributeLayoutAttr getLayoutBAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_b); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutB(); ::mlir::xegpu::DistributeLayoutAttr getLayoutCdAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_cd); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutCd(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAScaleAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_a_scale); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutAScale(); ::mlir::xegpu::DistributeLayoutAttr getLayoutBScaleAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_b_scale); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutBScale(); }; } // namespace detail template class DpasMxOpGenericAdaptor : public detail::DpasMxOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::DpasMxOpGenericAdaptorBase; static constexpr int odsIndex_a = 0; static constexpr int odsIndex_b = 1; static constexpr int odsIndex_acc = 2; static constexpr int odsIndex_scale_a = 3; static constexpr int odsIndex_scale_b = 4; public: DpasMxOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} DpasMxOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : DpasMxOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} DpasMxOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs) : DpasMxOpGenericAdaptor(values, attrs, Properties{}, {}) {} DpasMxOpGenericAdaptor(RangeT values, const DpasMxOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> DpasMxOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getA() { return (*getODSOperands(0).begin()); } ValueT getB() { return (*getODSOperands(1).begin()); } ValueT getAcc() { auto operands = getODSOperands(2); return operands.empty() ? ValueT{} : (*operands.begin()); } ValueT getScaleA() { auto operands = getODSOperands(3); return operands.empty() ? ValueT{} : (*operands.begin()); } ValueT getScaleB() { auto operands = getODSOperands(4); return operands.empty() ? ValueT{} : (*operands.begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class DpasMxOpAdaptor : public DpasMxOpGenericAdaptor<::mlir::ValueRange> { public: using DpasMxOpGenericAdaptor::DpasMxOpGenericAdaptor; DpasMxOpAdaptor(DpasMxOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class DpasMxOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::AtLeastNOperands<2>::Impl, ::mlir::OpTrait::AttrSizedOperandSegments, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::ConditionallySpeculatable::Trait, ::mlir::OpTrait::AlwaysSpeculatableImplTrait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = DpasMxOpAdaptor; template using GenericAdaptor = DpasMxOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_a = 0; static constexpr int odsIndex_b = 1; static constexpr int odsIndex_acc = 2; static constexpr int odsIndex_scale_a = 3; static constexpr int odsIndex_scale_b = 4; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("layout_a"), ::llvm::StringRef("layout_a_scale"), ::llvm::StringRef("layout_b"), ::llvm::StringRef("layout_b_scale"), ::llvm::StringRef("layout_cd"), ::llvm::StringRef("operandSegmentSizes")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getLayoutAAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getLayoutAAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getLayoutAScaleAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getLayoutAScaleAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getLayoutBAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getLayoutBAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } ::mlir::StringAttr getLayoutBScaleAttrName() { return getAttributeNameForIndex(3); } static ::mlir::StringAttr getLayoutBScaleAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 3); } ::mlir::StringAttr getLayoutCdAttrName() { return getAttributeNameForIndex(4); } static ::mlir::StringAttr getLayoutCdAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 4); } ::mlir::StringAttr getOperandSegmentSizesAttrName() { return (*this)->getName().getAttributeNames().back(); } static ::mlir::StringAttr getOperandSegmentSizesAttrName(::mlir::OperationName name) { return name.getAttributeNames().back(); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.dpas_mx"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getA() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::VectorType> getB() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(1).begin()); } ::mlir::TypedValue<::mlir::VectorType> getAcc() { auto operands = getODSOperands(2); return operands.empty() ? ::mlir::TypedValue<::mlir::VectorType>{} : ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*operands.begin()); } ::mlir::TypedValue<::mlir::Type> getScaleA() { auto operands = getODSOperands(3); return operands.empty() ? ::mlir::TypedValue<::mlir::Type>{} : ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*operands.begin()); } ::mlir::TypedValue<::mlir::Type> getScaleB() { auto operands = getODSOperands(4); return operands.empty() ? ::mlir::TypedValue<::mlir::Type>{} : ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*operands.begin()); } ::mlir::OpOperand &getAMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getBMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getAccMutable(); ::mlir::MutableOperandRange getScaleAMutable(); ::mlir::MutableOperandRange getScaleBMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getResult() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::xegpu::DistributeLayoutAttr getLayoutAAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_a); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutA(); ::mlir::xegpu::DistributeLayoutAttr getLayoutBAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_b); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutB(); ::mlir::xegpu::DistributeLayoutAttr getLayoutCdAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_cd); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutCd(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAScaleAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_a_scale); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutAScale(); ::mlir::xegpu::DistributeLayoutAttr getLayoutBScaleAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_b_scale); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutBScale(); void setLayoutAAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_a = attr; } void setLayoutBAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_b = attr; } void setLayoutCdAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_cd = attr; } void setLayoutAScaleAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_a_scale = attr; } void setLayoutBScaleAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_b_scale = attr; } ::mlir::Attribute removeLayoutAAttr() { auto attr = getProperties().layout_a; getProperties().layout_a = {}; return attr; } ::mlir::Attribute removeLayoutBAttr() { auto attr = getProperties().layout_b; getProperties().layout_b = {}; return attr; } ::mlir::Attribute removeLayoutCdAttr() { auto attr = getProperties().layout_cd; getProperties().layout_cd = {}; return attr; } ::mlir::Attribute removeLayoutAScaleAttr() { auto attr = getProperties().layout_a_scale; getProperties().layout_a_scale = {}; return attr; } ::mlir::Attribute removeLayoutBScaleAttr() { auto attr = getProperties().layout_b_scale; getProperties().layout_b_scale = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type result, ::mlir::Value a, ::mlir::Value b, /*optional*/::mlir::Value acc, /*optional*/::mlir::Value scale_a, /*optional*/::mlir::Value scale_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a_scale, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b_scale); static DpasMxOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type result, ::mlir::Value a, ::mlir::Value b, /*optional*/::mlir::Value acc, /*optional*/::mlir::Value scale_a, /*optional*/::mlir::Value scale_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a_scale, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b_scale); static DpasMxOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type result, ::mlir::Value a, ::mlir::Value b, /*optional*/::mlir::Value acc, /*optional*/::mlir::Value scale_a, /*optional*/::mlir::Value scale_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a_scale, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b_scale); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, /*optional*/::mlir::Value acc, /*optional*/::mlir::Value scale_a, /*optional*/::mlir::Value scale_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a_scale, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b_scale); static DpasMxOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, /*optional*/::mlir::Value acc, /*optional*/::mlir::Value scale_a, /*optional*/::mlir::Value scale_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a_scale, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b_scale); static DpasMxOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, /*optional*/::mlir::Value acc, /*optional*/::mlir::Value scale_a, /*optional*/::mlir::Value scale_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a_scale, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b_scale); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static DpasMxOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static DpasMxOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static DpasMxOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static DpasMxOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 5 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayoutCd().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutCdAttr(anchorLayout); } VectorType getAType() { return getA().getType(); } VectorType getBType() { return getB().getType(); } VectorType getAccType() { return getAcc().getType(); } Type getScaleAType() { return getScaleA().getType(); } Type getScaleBType() { return getScaleB().getType(); } VectorType getResultType() { return getResult().getType(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::DpasMxOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::DpasMxOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::DpasOp declarations //===----------------------------------------------------------------------===// namespace detail { class DpasOpGenericAdaptorBase { public: struct Properties { using layout_aTy = ::mlir::xegpu::DistributeLayoutAttr; layout_aTy layout_a; auto getLayoutA() const { auto &propStorage = this->layout_a; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutA(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_a = propValue; } using layout_bTy = ::mlir::xegpu::DistributeLayoutAttr; layout_bTy layout_b; auto getLayoutB() const { auto &propStorage = this->layout_b; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutB(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_b = propValue; } using layout_cdTy = ::mlir::xegpu::DistributeLayoutAttr; layout_cdTy layout_cd; auto getLayoutCd() const { auto &propStorage = this->layout_cd; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayoutCd(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout_cd = propValue; } bool operator==(const Properties &rhs) const { return rhs.layout_a == this->layout_a && rhs.layout_b == this->layout_b && rhs.layout_cd == this->layout_cd && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: DpasOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.dpas", odsAttrs.getContext()); } DpasOpGenericAdaptorBase(DpasOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::xegpu::DistributeLayoutAttr getLayoutAAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_a); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutA(); ::mlir::xegpu::DistributeLayoutAttr getLayoutBAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_b); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutB(); ::mlir::xegpu::DistributeLayoutAttr getLayoutCdAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_cd); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutCd(); }; } // namespace detail template class DpasOpGenericAdaptor : public detail::DpasOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::DpasOpGenericAdaptorBase; static constexpr int odsIndex_lhs = 0; static constexpr int odsIndex_rhs = 1; static constexpr int odsIndex_acc = 2; public: DpasOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} DpasOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : DpasOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} DpasOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : DpasOpGenericAdaptor(values, attrs, Properties{}, {}) {} DpasOpGenericAdaptor(RangeT values, const DpasOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> DpasOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getLhs() { return (*getODSOperands(0).begin()); } ValueT getRhs() { return (*getODSOperands(1).begin()); } ValueT getAcc() { auto operands = getODSOperands(2); return operands.empty() ? ValueT{} : (*operands.begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class DpasOpAdaptor : public DpasOpGenericAdaptor<::mlir::ValueRange> { public: using DpasOpGenericAdaptor::DpasOpGenericAdaptor; DpasOpAdaptor(DpasOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class DpasOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::AtLeastNOperands<2>::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::ConditionallySpeculatable::Trait, ::mlir::OpTrait::AlwaysSpeculatableImplTrait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = DpasOpAdaptor; template using GenericAdaptor = DpasOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_lhs = 0; static constexpr int odsIndex_rhs = 1; static constexpr int odsIndex_acc = 2; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("layout_a"), ::llvm::StringRef("layout_b"), ::llvm::StringRef("layout_cd")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getLayoutAAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getLayoutAAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getLayoutBAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getLayoutBAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getLayoutCdAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getLayoutCdAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.dpas"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getLhs() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::VectorType> getRhs() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(1).begin()); } ::mlir::TypedValue<::mlir::VectorType> getAcc() { auto operands = getODSOperands(2); return operands.empty() ? ::mlir::TypedValue<::mlir::VectorType>{} : ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*operands.begin()); } ::mlir::OpOperand &getLhsMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getRhsMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getAccMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getResult() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::xegpu::DistributeLayoutAttr getLayoutAAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_a); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutA(); ::mlir::xegpu::DistributeLayoutAttr getLayoutBAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_b); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutB(); ::mlir::xegpu::DistributeLayoutAttr getLayoutCdAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout_cd); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayoutCd(); void setLayoutAAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_a = attr; } void setLayoutBAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_b = attr; } void setLayoutCdAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout_cd = attr; } ::mlir::Attribute removeLayoutAAttr() { auto attr = getProperties().layout_a; getProperties().layout_a = {}; return attr; } ::mlir::Attribute removeLayoutBAttr() { auto attr = getProperties().layout_b; getProperties().layout_b = {}; return attr; } ::mlir::Attribute removeLayoutCdAttr() { auto attr = getProperties().layout_cd; getProperties().layout_cd = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type result, ::mlir::Value lhs, ::mlir::Value rhs, /*optional*/::mlir::Value acc, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd); static DpasOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type result, ::mlir::Value lhs, ::mlir::Value rhs, /*optional*/::mlir::Value acc, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd); static DpasOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type result, ::mlir::Value lhs, ::mlir::Value rhs, /*optional*/::mlir::Value acc, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value lhs, ::mlir::Value rhs, /*optional*/::mlir::Value acc, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd); static DpasOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value lhs, ::mlir::Value rhs, /*optional*/::mlir::Value acc, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd); static DpasOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value lhs, ::mlir::Value rhs, /*optional*/::mlir::Value acc, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_a, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_b, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout_cd); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static DpasOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static DpasOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static DpasOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static DpasOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 3 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayoutCd().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutCdAttr(anchorLayout); } VectorType getLhsType() { return getLhs().getType(); } VectorType getRhsType() { return getRhs().getType(); } VectorType getAccType() { if (getAcc()) return getAcc().getType(); return {}; } VectorType getResultType() { return getResult().getType(); } bool hasAcc() { return getAcc() != nullptr; } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::DpasOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::DpasOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::FenceOp declarations //===----------------------------------------------------------------------===// namespace detail { class FenceOpGenericAdaptorBase { public: struct Properties { using fence_scopeTy = ::mlir::xegpu::FenceScopeAttr; fence_scopeTy fence_scope; auto getFenceScope() const { auto &propStorage = this->fence_scope; return ::llvm::cast<::mlir::xegpu::FenceScopeAttr>(propStorage); } void setFenceScope(const ::mlir::xegpu::FenceScopeAttr &propValue) { this->fence_scope = propValue; } using memory_kindTy = ::mlir::xegpu::MemorySpaceAttr; memory_kindTy memory_kind; auto getMemoryKind() const { auto &propStorage = this->memory_kind; return ::llvm::cast<::mlir::xegpu::MemorySpaceAttr>(propStorage); } void setMemoryKind(const ::mlir::xegpu::MemorySpaceAttr &propValue) { this->memory_kind = propValue; } bool operator==(const Properties &rhs) const { return rhs.fence_scope == this->fence_scope && rhs.memory_kind == this->memory_kind && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: FenceOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.fence", odsAttrs.getContext()); } FenceOpGenericAdaptorBase(FenceOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::xegpu::MemorySpaceAttr getMemoryKindAttr() { auto attr = ::llvm::cast<::mlir::xegpu::MemorySpaceAttr>(getProperties().memory_kind); return attr; } ::mlir::xegpu::MemorySpace getMemoryKind(); ::mlir::xegpu::FenceScopeAttr getFenceScopeAttr() { auto attr = ::llvm::cast<::mlir::xegpu::FenceScopeAttr>(getProperties().fence_scope); return attr; } ::mlir::xegpu::FenceScope getFenceScope(); }; } // namespace detail template class FenceOpGenericAdaptor : public detail::FenceOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::FenceOpGenericAdaptorBase; public: FenceOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} FenceOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : FenceOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} FenceOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : FenceOpGenericAdaptor(values, attrs, Properties{}, {}) {} FenceOpGenericAdaptor(RangeT values, const FenceOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> FenceOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class FenceOpAdaptor : public FenceOpGenericAdaptor<::mlir::ValueRange> { public: using FenceOpGenericAdaptor::FenceOpGenericAdaptor; FenceOpAdaptor(FenceOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class FenceOp : public ::mlir::Op { public: using Op::Op; using Op::print; using Adaptor = FenceOpAdaptor; template using GenericAdaptor = FenceOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("fence_scope"), ::llvm::StringRef("memory_kind")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getFenceScopeAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getFenceScopeAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getMemoryKindAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getMemoryKindAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.fence"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::xegpu::MemorySpaceAttr getMemoryKindAttr() { return ::llvm::cast<::mlir::xegpu::MemorySpaceAttr>(getProperties().memory_kind); } ::mlir::xegpu::MemorySpace getMemoryKind(); ::mlir::xegpu::FenceScopeAttr getFenceScopeAttr() { return ::llvm::cast<::mlir::xegpu::FenceScopeAttr>(getProperties().fence_scope); } ::mlir::xegpu::FenceScope getFenceScope(); void setMemoryKindAttr(::mlir::xegpu::MemorySpaceAttr attr) { getProperties().memory_kind = attr; } void setMemoryKind(::mlir::xegpu::MemorySpace attrValue); void setFenceScopeAttr(::mlir::xegpu::FenceScopeAttr attr) { getProperties().fence_scope = attr; } void setFenceScope(::mlir::xegpu::FenceScope attrValue); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::xegpu::MemorySpaceAttr memory_kind, ::mlir::xegpu::FenceScopeAttr fence_scope); static FenceOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::xegpu::MemorySpaceAttr memory_kind, ::mlir::xegpu::FenceScopeAttr fence_scope); static FenceOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::xegpu::MemorySpaceAttr memory_kind, ::mlir::xegpu::FenceScopeAttr fence_scope); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::xegpu::MemorySpaceAttr memory_kind, ::mlir::xegpu::FenceScopeAttr fence_scope); static FenceOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::xegpu::MemorySpaceAttr memory_kind, ::mlir::xegpu::FenceScopeAttr fence_scope); static FenceOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::xegpu::MemorySpaceAttr memory_kind, ::mlir::xegpu::FenceScopeAttr fence_scope); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::xegpu::MemorySpace memory_kind, ::mlir::xegpu::FenceScope fence_scope); static FenceOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::xegpu::MemorySpace memory_kind, ::mlir::xegpu::FenceScope fence_scope); static FenceOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::xegpu::MemorySpace memory_kind, ::mlir::xegpu::FenceScope fence_scope); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::xegpu::MemorySpace memory_kind, ::mlir::xegpu::FenceScope fence_scope); static FenceOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::xegpu::MemorySpace memory_kind, ::mlir::xegpu::FenceScope fence_scope); static FenceOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::xegpu::MemorySpace memory_kind, ::mlir::xegpu::FenceScope fence_scope); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static FenceOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static FenceOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static FenceOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static FenceOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 2 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::FenceOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::FenceOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::InitNbarrierOp declarations //===----------------------------------------------------------------------===// namespace detail { class InitNbarrierOpGenericAdaptorBase { public: using Properties = ::mlir::EmptyProperties; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: InitNbarrierOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.init_nbarrier", odsAttrs.getContext()); } InitNbarrierOpGenericAdaptorBase(InitNbarrierOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } }; } // namespace detail template class InitNbarrierOpGenericAdaptor : public detail::InitNbarrierOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::InitNbarrierOpGenericAdaptorBase; static constexpr int odsIndex_nbarrier_id = 0; static constexpr int odsIndex_participant_thread_num = 1; public: InitNbarrierOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} InitNbarrierOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : InitNbarrierOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} InitNbarrierOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : InitNbarrierOpGenericAdaptor(values, attrs, Properties{}, {}) {} InitNbarrierOpGenericAdaptor(RangeT values, const InitNbarrierOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> InitNbarrierOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getNbarrierId() { return (*getODSOperands(0).begin()); } ValueT getParticipantThreadNum() { return (*getODSOperands(1).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class InitNbarrierOpAdaptor : public InitNbarrierOpGenericAdaptor<::mlir::ValueRange> { public: using InitNbarrierOpGenericAdaptor::InitNbarrierOpGenericAdaptor; InitNbarrierOpAdaptor(InitNbarrierOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class InitNbarrierOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::NOperands<2>::Impl, ::mlir::OpTrait::OpInvariants> { public: using Op::Op; using Op::print; using Adaptor = InitNbarrierOpAdaptor; template using GenericAdaptor = InitNbarrierOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_nbarrier_id = 0; static constexpr int odsIndex_participant_thread_num = 1; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { return {}; } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.init_nbarrier"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::IntegerType> getNbarrierId() { return ::llvm::cast<::mlir::TypedValue<::mlir::IntegerType>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::IntegerType> getParticipantThreadNum() { return ::llvm::cast<::mlir::TypedValue<::mlir::IntegerType>>(*getODSOperands(1).begin()); } ::mlir::OpOperand &getNbarrierIdMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getParticipantThreadNumMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::NbarrierType> getResult() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::NbarrierType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type result, ::mlir::Value nbarrier_id, ::mlir::Value participant_thread_num); static InitNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type result, ::mlir::Value nbarrier_id, ::mlir::Value participant_thread_num); static InitNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type result, ::mlir::Value nbarrier_id, ::mlir::Value participant_thread_num); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier_id, ::mlir::Value participant_thread_num); static InitNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier_id, ::mlir::Value participant_thread_num); static InitNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier_id, ::mlir::Value participant_thread_num); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static InitNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static InitNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static InitNbarrierOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static InitNbarrierOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); public: }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::InitNbarrierOp) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::LoadGatherOp declarations //===----------------------------------------------------------------------===// namespace detail { class LoadGatherOpGenericAdaptorBase { public: struct Properties { using chunk_sizeTy = ::mlir::IntegerAttr; chunk_sizeTy chunk_size; auto getChunkSize() const { auto &propStorage = this->chunk_size; return ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(propStorage); } void setChunkSize(const ::mlir::IntegerAttr &propValue) { this->chunk_size = propValue; } using l1_hintTy = ::mlir::xegpu::CachePolicyAttr; l1_hintTy l1_hint; auto getL1Hint() const { auto &propStorage = this->l1_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL1Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l1_hint = propValue; } using l2_hintTy = ::mlir::xegpu::CachePolicyAttr; l2_hintTy l2_hint; auto getL2Hint() const { auto &propStorage = this->l2_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL2Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l2_hint = propValue; } using l3_hintTy = ::mlir::xegpu::CachePolicyAttr; l3_hintTy l3_hint; auto getL3Hint() const { auto &propStorage = this->l3_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL3Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l3_hint = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } bool operator==(const Properties &rhs) const { return rhs.chunk_size == this->chunk_size && rhs.l1_hint == this->l1_hint && rhs.l2_hint == this->l2_hint && rhs.l3_hint == this->l3_hint && rhs.layout == this->layout && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: LoadGatherOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.load", odsAttrs.getContext()); } LoadGatherOpGenericAdaptorBase(LoadGatherOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::IntegerAttr getChunkSizeAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(getProperties().chunk_size); return attr; } ::std::optional getChunkSize(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class LoadGatherOpGenericAdaptor : public detail::LoadGatherOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::LoadGatherOpGenericAdaptorBase; static constexpr int odsIndex_source = 0; static constexpr int odsIndex_offsets = 1; static constexpr int odsIndex_mask = 2; public: LoadGatherOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} LoadGatherOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : LoadGatherOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} LoadGatherOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : LoadGatherOpGenericAdaptor(values, attrs, Properties{}, {}) {} LoadGatherOpGenericAdaptor(RangeT values, const LoadGatherOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> LoadGatherOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getSource() { return (*getODSOperands(0).begin()); } ValueT getOffsets() { return (*getODSOperands(1).begin()); } ValueT getMask() { return (*getODSOperands(2).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class LoadGatherOpAdaptor : public LoadGatherOpGenericAdaptor<::mlir::ValueRange> { public: using LoadGatherOpGenericAdaptor::LoadGatherOpGenericAdaptor; LoadGatherOpAdaptor(LoadGatherOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class LoadGatherOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::NOperands<3>::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = LoadGatherOpAdaptor; template using GenericAdaptor = LoadGatherOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_source = 0; static constexpr int odsIndex_offsets = 1; static constexpr int odsIndex_mask = 2; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("chunk_size"), ::llvm::StringRef("l1_hint"), ::llvm::StringRef("l2_hint"), ::llvm::StringRef("l3_hint"), ::llvm::StringRef("layout")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getChunkSizeAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getChunkSizeAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getL1HintAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getL1HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getL2HintAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getL2HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } ::mlir::StringAttr getL3HintAttrName() { return getAttributeNameForIndex(3); } static ::mlir::StringAttr getL3HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 3); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(4); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 4); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.load"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getSource() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::Type> getOffsets() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(1).begin()); } ::mlir::TypedValue<::mlir::Type> getMask() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(2).begin()); } ::mlir::OpOperand &getSourceMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getOffsetsMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getMaskMutable() { auto range = getODSOperandIndexAndLength(2); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getValue() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::IntegerAttr getChunkSizeAttr() { return ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(getProperties().chunk_size); } ::std::optional getChunkSize(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setChunkSizeAttr(::mlir::IntegerAttr attr) { getProperties().chunk_size = attr; } void setChunkSize(::std::optional attrValue); void setL1HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l1_hint = attr; } void setL1Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL2HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l2_hint = attr; } void setL2Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL3HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l3_hint = attr; } void setL3Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeChunkSizeAttr() { auto attr = getProperties().chunk_size; getProperties().chunk_size = {}; return attr; } ::mlir::Attribute removeL1HintAttr() { auto attr = getProperties().l1_hint; getProperties().l1_hint = {}; return attr; } ::mlir::Attribute removeL2HintAttr() { auto attr = getProperties().l2_hint; getProperties().l2_hint = {}; return attr; } ::mlir::Attribute removeL3HintAttr() { auto attr = getProperties().l3_hint; getProperties().l3_hint = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Type value, Value source, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint); static LoadGatherOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Type value, Value source, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint); static LoadGatherOp create(::mlir::ImplicitLocOpBuilder &builder, Type value, Value source, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Type value, Value source, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static LoadGatherOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Type value, Value source, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static LoadGatherOp create(::mlir::ImplicitLocOpBuilder &builder, Type value, Value source, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type value, ::mlir::Value source, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadGatherOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type value, ::mlir::Value source, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadGatherOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type value, ::mlir::Value source, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadGatherOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadGatherOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static LoadGatherOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static LoadGatherOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static LoadGatherOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static LoadGatherOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 5 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } Type getSourceType() { return getSource().getType(); } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } mlir::Type getElementType() { auto type = getValue().getType(); return getElementTypeOrSelf(type); } VectorType getValueType() { return llvm::dyn_cast(getValue().getType()); } Type getMaskType() { return getMask().getType(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::LoadGatherOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::LoadGatherOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::LoadMatrixOp declarations //===----------------------------------------------------------------------===// namespace detail { class LoadMatrixOpGenericAdaptorBase { public: struct Properties { using const_offsetsTy = ::mlir::DenseI64ArrayAttr; const_offsetsTy const_offsets; auto getConstOffsets() const { auto &propStorage = this->const_offsets; return ::llvm::cast<::mlir::DenseI64ArrayAttr>(propStorage); } void setConstOffsets(const ::mlir::DenseI64ArrayAttr &propValue) { this->const_offsets = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } using subgroup_block_ioTy = ::mlir::UnitAttr; subgroup_block_ioTy subgroup_block_io; auto getSubgroupBlockIo() const { auto &propStorage = this->subgroup_block_io; return ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(propStorage); } void setSubgroupBlockIo(const ::mlir::UnitAttr &propValue) { this->subgroup_block_io = propValue; } bool operator==(const Properties &rhs) const { return rhs.const_offsets == this->const_offsets && rhs.layout == this->layout && rhs.subgroup_block_io == this->subgroup_block_io && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: LoadMatrixOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.load_matrix", odsAttrs.getContext()); } LoadMatrixOpGenericAdaptorBase(LoadMatrixOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { auto attr = ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); return attr; } ::llvm::ArrayRef getConstOffsets(); ::mlir::UnitAttr getSubgroupBlockIoAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(getProperties().subgroup_block_io); return attr; } ::std::optional getSubgroupBlockIo(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class LoadMatrixOpGenericAdaptor : public detail::LoadMatrixOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::LoadMatrixOpGenericAdaptorBase; static constexpr int odsIndex_mem_desc = 0; static constexpr int odsIndex_offsets = 1; public: LoadMatrixOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} LoadMatrixOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : LoadMatrixOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} LoadMatrixOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : LoadMatrixOpGenericAdaptor(values, attrs, Properties{}, {}) {} LoadMatrixOpGenericAdaptor(RangeT values, const LoadMatrixOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> LoadMatrixOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getMemDesc() { return (*getODSOperands(0).begin()); } RangeT getOffsets() { return getODSOperands(1); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class LoadMatrixOpAdaptor : public LoadMatrixOpGenericAdaptor<::mlir::ValueRange> { public: using LoadMatrixOpGenericAdaptor::LoadMatrixOpGenericAdaptor; LoadMatrixOpAdaptor(LoadMatrixOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class LoadMatrixOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::AtLeastNOperands<1>::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = LoadMatrixOpAdaptor; template using GenericAdaptor = LoadMatrixOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_mem_desc = 0; static constexpr int odsIndex_offsets = 1; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("const_offsets"), ::llvm::StringRef("layout"), ::llvm::StringRef("subgroup_block_io")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getConstOffsetsAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getConstOffsetsAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getSubgroupBlockIoAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getSubgroupBlockIoAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.load_matrix"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::MemDescType> getMemDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::MemDescType>>(*getODSOperands(0).begin()); } ::mlir::Operation::operand_range getOffsets() { return getODSOperands(1); } ::mlir::OpOperand &getMemDescMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getOffsetsMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getRes() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { return ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); } ::llvm::ArrayRef getConstOffsets(); ::mlir::UnitAttr getSubgroupBlockIoAttr() { return ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(getProperties().subgroup_block_io); } ::std::optional getSubgroupBlockIo(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setConstOffsetsAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().const_offsets = attr; } void setConstOffsets(::llvm::ArrayRef attrValue); void setSubgroupBlockIoAttr(::mlir::UnitAttr attr) { getProperties().subgroup_block_io = attr; } void setSubgroupBlockIo(bool attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeSubgroupBlockIoAttr() { auto attr = getProperties().subgroup_block_io; getProperties().subgroup_block_io = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Type res, TypedValue mem_desc, llvm::ArrayRef offsets, DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Type res, TypedValue mem_desc, llvm::ArrayRef offsets, DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, Type res, TypedValue mem_desc, llvm::ArrayRef offsets, DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type res, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type res, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type res, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type res, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type res, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type res, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static LoadMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static LoadMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static LoadMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static LoadMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 3 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: SmallVector getMixedOffsets() { return getMixedValues(getConstOffsets(), getOffsets(), getContext()); } ArrayRef getDataShape() { auto resTy = getRes().getType(); if (auto vecTy = llvm::dyn_cast(resTy)) return vecTy.getShape(); return {}; } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::LoadMatrixOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::LoadMatrixOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::LoadNdOp declarations //===----------------------------------------------------------------------===// namespace detail { class LoadNdOpGenericAdaptorBase { public: struct Properties { using const_offsetsTy = ::mlir::DenseI64ArrayAttr; const_offsetsTy const_offsets; auto getConstOffsets() const { auto &propStorage = this->const_offsets; return ::llvm::cast<::mlir::DenseI64ArrayAttr>(propStorage); } void setConstOffsets(const ::mlir::DenseI64ArrayAttr &propValue) { this->const_offsets = propValue; } using l1_hintTy = ::mlir::xegpu::CachePolicyAttr; l1_hintTy l1_hint; auto getL1Hint() const { auto &propStorage = this->l1_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL1Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l1_hint = propValue; } using l2_hintTy = ::mlir::xegpu::CachePolicyAttr; l2_hintTy l2_hint; auto getL2Hint() const { auto &propStorage = this->l2_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL2Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l2_hint = propValue; } using l3_hintTy = ::mlir::xegpu::CachePolicyAttr; l3_hintTy l3_hint; auto getL3Hint() const { auto &propStorage = this->l3_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL3Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l3_hint = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } using packedTy = ::mlir::UnitAttr; packedTy packed; auto getPacked() const { auto &propStorage = this->packed; return ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(propStorage); } void setPacked(const ::mlir::UnitAttr &propValue) { this->packed = propValue; } using transposeTy = ::mlir::DenseI64ArrayAttr; transposeTy transpose; auto getTranspose() const { auto &propStorage = this->transpose; return ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(propStorage); } void setTranspose(const ::mlir::DenseI64ArrayAttr &propValue) { this->transpose = propValue; } bool operator==(const Properties &rhs) const { return rhs.const_offsets == this->const_offsets && rhs.l1_hint == this->l1_hint && rhs.l2_hint == this->l2_hint && rhs.l3_hint == this->l3_hint && rhs.layout == this->layout && rhs.packed == this->packed && rhs.transpose == this->transpose && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: LoadNdOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.load_nd", odsAttrs.getContext()); } LoadNdOpGenericAdaptorBase(LoadNdOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { auto attr = ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); return attr; } ::llvm::ArrayRef getConstOffsets(); ::mlir::UnitAttr getPackedAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(getProperties().packed); return attr; } ::std::optional getPacked(); ::mlir::DenseI64ArrayAttr getTransposeAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(getProperties().transpose); return attr; } ::std::optional<::llvm::ArrayRef> getTranspose(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class LoadNdOpGenericAdaptor : public detail::LoadNdOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::LoadNdOpGenericAdaptorBase; static constexpr int odsIndex_TensorDesc = 0; static constexpr int odsIndex_offsets = 1; public: LoadNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} LoadNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : LoadNdOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} LoadNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : LoadNdOpGenericAdaptor(values, attrs, Properties{}, {}) {} LoadNdOpGenericAdaptor(RangeT values, const LoadNdOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> LoadNdOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getTensorDesc() { return (*getODSOperands(0).begin()); } RangeT getOffsets() { return getODSOperands(1); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class LoadNdOpAdaptor : public LoadNdOpGenericAdaptor<::mlir::ValueRange> { public: using LoadNdOpGenericAdaptor::LoadNdOpGenericAdaptor; LoadNdOpAdaptor(LoadNdOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class LoadNdOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::AtLeastNOperands<1>::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = LoadNdOpAdaptor; template using GenericAdaptor = LoadNdOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_TensorDesc = 0; static constexpr int odsIndex_offsets = 1; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("const_offsets"), ::llvm::StringRef("l1_hint"), ::llvm::StringRef("l2_hint"), ::llvm::StringRef("l3_hint"), ::llvm::StringRef("layout"), ::llvm::StringRef("packed"), ::llvm::StringRef("transpose")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getConstOffsetsAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getConstOffsetsAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getL1HintAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getL1HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getL2HintAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getL2HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } ::mlir::StringAttr getL3HintAttrName() { return getAttributeNameForIndex(3); } static ::mlir::StringAttr getL3HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 3); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(4); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 4); } ::mlir::StringAttr getPackedAttrName() { return getAttributeNameForIndex(5); } static ::mlir::StringAttr getPackedAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 5); } ::mlir::StringAttr getTransposeAttrName() { return getAttributeNameForIndex(6); } static ::mlir::StringAttr getTransposeAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 6); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.load_nd"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::TensorDescType> getTensorDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::TensorDescType>>(*getODSOperands(0).begin()); } ::mlir::Operation::operand_range getOffsets() { return getODSOperands(1); } ::mlir::OpOperand &getTensorDescMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getOffsetsMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getValue() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { return ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); } ::llvm::ArrayRef getConstOffsets(); ::mlir::UnitAttr getPackedAttr() { return ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(getProperties().packed); } ::std::optional getPacked(); ::mlir::DenseI64ArrayAttr getTransposeAttr() { return ::llvm::dyn_cast_or_null<::mlir::DenseI64ArrayAttr>(getProperties().transpose); } ::std::optional<::llvm::ArrayRef> getTranspose(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setConstOffsetsAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().const_offsets = attr; } void setConstOffsets(::llvm::ArrayRef attrValue); void setPackedAttr(::mlir::UnitAttr attr) { getProperties().packed = attr; } void setPacked(bool attrValue); void setTransposeAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().transpose = attr; } void setTranspose(::std::optional<::llvm::ArrayRef> attrValue); void setL1HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l1_hint = attr; } void setL1Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL2HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l2_hint = attr; } void setL2Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL3HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l3_hint = attr; } void setL3Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removePackedAttr() { auto attr = getProperties().packed; getProperties().packed = {}; return attr; } ::mlir::Attribute removeTransposeAttr() { auto attr = getProperties().transpose; getProperties().transpose = {}; return attr; } ::mlir::Attribute removeL1HintAttr() { auto attr = getProperties().l1_hint; getProperties().l1_hint = {}; return attr; } ::mlir::Attribute removeL2HintAttr() { auto attr = getProperties().l2_hint; getProperties().l2_hint = {}; return attr; } ::mlir::Attribute removeL3HintAttr() { auto attr = getProperties().l3_hint; getProperties().l3_hint = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Type value, Value TensorDesc, ArrayRef offsets, UnitAttr packed, DenseI64ArrayAttr transpose, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Type value, Value TensorDesc, ArrayRef offsets, UnitAttr packed, DenseI64ArrayAttr transpose, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::ImplicitLocOpBuilder &builder, Type value, Value TensorDesc, ArrayRef offsets, UnitAttr packed, DenseI64ArrayAttr transpose, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static LoadNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr packed, /*optional*/::mlir::DenseI64ArrayAttr transpose, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static LoadNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static LoadNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static LoadNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static LoadNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 7 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } VectorType getType() { return llvm::dyn_cast(getValue().getType()); } xegpu::TensorDescType getTensorDescType() { return getTensorDesc().getType(); } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } SmallVector getMixedOffsets() { return getMixedValues(getConstOffsets(), getOffsets(), getContext()); } xegpu::DistributeLayoutAttr getDescLayoutAttr() { return dyn_cast_if_present(getTensorDescType().getLayout()); } ArrayRef getDataShape() { return getTensorDescType().getShape(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::LoadNdOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::LoadNdOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::NbarrierArriveOp declarations //===----------------------------------------------------------------------===// namespace detail { class NbarrierArriveOpGenericAdaptorBase { public: using Properties = ::mlir::EmptyProperties; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: NbarrierArriveOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.nbarrier_arrive", odsAttrs.getContext()); } NbarrierArriveOpGenericAdaptorBase(NbarrierArriveOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } }; } // namespace detail template class NbarrierArriveOpGenericAdaptor : public detail::NbarrierArriveOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::NbarrierArriveOpGenericAdaptorBase; static constexpr int odsIndex_nbarrier = 0; public: NbarrierArriveOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} NbarrierArriveOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : NbarrierArriveOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} NbarrierArriveOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : NbarrierArriveOpGenericAdaptor(values, attrs, Properties{}, {}) {} NbarrierArriveOpGenericAdaptor(RangeT values, const NbarrierArriveOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> NbarrierArriveOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getNbarrier() { return (*getODSOperands(0).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class NbarrierArriveOpAdaptor : public NbarrierArriveOpGenericAdaptor<::mlir::ValueRange> { public: using NbarrierArriveOpGenericAdaptor::NbarrierArriveOpGenericAdaptor; NbarrierArriveOpAdaptor(NbarrierArriveOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class NbarrierArriveOp : public ::mlir::Op { public: using Op::Op; using Op::print; using Adaptor = NbarrierArriveOpAdaptor; template using GenericAdaptor = NbarrierArriveOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_nbarrier = 0; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { return {}; } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.nbarrier_arrive"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::NbarrierType> getNbarrier() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::NbarrierType>>(*getODSOperands(0).begin()); } ::mlir::OpOperand &getNbarrierMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value nbarrier); static NbarrierArriveOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value nbarrier); static NbarrierArriveOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value nbarrier); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier); static NbarrierArriveOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier); static NbarrierArriveOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static NbarrierArriveOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static NbarrierArriveOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static NbarrierArriveOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static NbarrierArriveOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); public: }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::NbarrierArriveOp) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::NbarrierWaitOp declarations //===----------------------------------------------------------------------===// namespace detail { class NbarrierWaitOpGenericAdaptorBase { public: using Properties = ::mlir::EmptyProperties; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: NbarrierWaitOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.nbarrier_wait", odsAttrs.getContext()); } NbarrierWaitOpGenericAdaptorBase(NbarrierWaitOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } }; } // namespace detail template class NbarrierWaitOpGenericAdaptor : public detail::NbarrierWaitOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::NbarrierWaitOpGenericAdaptorBase; static constexpr int odsIndex_nbarrier = 0; public: NbarrierWaitOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} NbarrierWaitOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : NbarrierWaitOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} NbarrierWaitOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : NbarrierWaitOpGenericAdaptor(values, attrs, Properties{}, {}) {} NbarrierWaitOpGenericAdaptor(RangeT values, const NbarrierWaitOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> NbarrierWaitOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getNbarrier() { return (*getODSOperands(0).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class NbarrierWaitOpAdaptor : public NbarrierWaitOpGenericAdaptor<::mlir::ValueRange> { public: using NbarrierWaitOpGenericAdaptor::NbarrierWaitOpGenericAdaptor; NbarrierWaitOpAdaptor(NbarrierWaitOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class NbarrierWaitOp : public ::mlir::Op { public: using Op::Op; using Op::print; using Adaptor = NbarrierWaitOpAdaptor; template using GenericAdaptor = NbarrierWaitOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_nbarrier = 0; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { return {}; } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.nbarrier_wait"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::NbarrierType> getNbarrier() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::NbarrierType>>(*getODSOperands(0).begin()); } ::mlir::OpOperand &getNbarrierMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value nbarrier); static NbarrierWaitOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value nbarrier); static NbarrierWaitOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value nbarrier); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier); static NbarrierWaitOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier); static NbarrierWaitOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value nbarrier); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static NbarrierWaitOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static NbarrierWaitOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static NbarrierWaitOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static NbarrierWaitOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); public: }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::NbarrierWaitOp) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::PrefetchNdOp declarations //===----------------------------------------------------------------------===// namespace detail { class PrefetchNdOpGenericAdaptorBase { public: struct Properties { using const_offsetsTy = ::mlir::DenseI64ArrayAttr; const_offsetsTy const_offsets; auto getConstOffsets() const { auto &propStorage = this->const_offsets; return ::llvm::cast<::mlir::DenseI64ArrayAttr>(propStorage); } void setConstOffsets(const ::mlir::DenseI64ArrayAttr &propValue) { this->const_offsets = propValue; } using l1_hintTy = ::mlir::xegpu::CachePolicyAttr; l1_hintTy l1_hint; auto getL1Hint() const { auto &propStorage = this->l1_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL1Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l1_hint = propValue; } using l2_hintTy = ::mlir::xegpu::CachePolicyAttr; l2_hintTy l2_hint; auto getL2Hint() const { auto &propStorage = this->l2_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL2Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l2_hint = propValue; } using l3_hintTy = ::mlir::xegpu::CachePolicyAttr; l3_hintTy l3_hint; auto getL3Hint() const { auto &propStorage = this->l3_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL3Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l3_hint = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } bool operator==(const Properties &rhs) const { return rhs.const_offsets == this->const_offsets && rhs.l1_hint == this->l1_hint && rhs.l2_hint == this->l2_hint && rhs.l3_hint == this->l3_hint && rhs.layout == this->layout && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: PrefetchNdOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.prefetch_nd", odsAttrs.getContext()); } PrefetchNdOpGenericAdaptorBase(PrefetchNdOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { auto attr = ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); return attr; } ::llvm::ArrayRef getConstOffsets(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class PrefetchNdOpGenericAdaptor : public detail::PrefetchNdOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::PrefetchNdOpGenericAdaptorBase; static constexpr int odsIndex_TensorDesc = 0; static constexpr int odsIndex_offsets = 1; public: PrefetchNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} PrefetchNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : PrefetchNdOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} PrefetchNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : PrefetchNdOpGenericAdaptor(values, attrs, Properties{}, {}) {} PrefetchNdOpGenericAdaptor(RangeT values, const PrefetchNdOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> PrefetchNdOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getTensorDesc() { return (*getODSOperands(0).begin()); } RangeT getOffsets() { return getODSOperands(1); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class PrefetchNdOpAdaptor : public PrefetchNdOpGenericAdaptor<::mlir::ValueRange> { public: using PrefetchNdOpGenericAdaptor::PrefetchNdOpGenericAdaptor; PrefetchNdOpAdaptor(PrefetchNdOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class PrefetchNdOp : public ::mlir::Op::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = PrefetchNdOpAdaptor; template using GenericAdaptor = PrefetchNdOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_TensorDesc = 0; static constexpr int odsIndex_offsets = 1; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("const_offsets"), ::llvm::StringRef("l1_hint"), ::llvm::StringRef("l2_hint"), ::llvm::StringRef("l3_hint"), ::llvm::StringRef("layout")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getConstOffsetsAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getConstOffsetsAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getL1HintAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getL1HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getL2HintAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getL2HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } ::mlir::StringAttr getL3HintAttrName() { return getAttributeNameForIndex(3); } static ::mlir::StringAttr getL3HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 3); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(4); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 4); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.prefetch_nd"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::xegpu::TensorDescType> getTensorDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::TensorDescType>>(*getODSOperands(0).begin()); } ::mlir::Operation::operand_range getOffsets() { return getODSOperands(1); } ::mlir::OpOperand &getTensorDescMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getOffsetsMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { return ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); } ::llvm::ArrayRef getConstOffsets(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setConstOffsetsAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().const_offsets = attr; } void setConstOffsets(::llvm::ArrayRef attrValue); void setL1HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l1_hint = attr; } void setL1Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL2HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l2_hint = attr; } void setL2Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL3HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l3_hint = attr; } void setL3Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeL1HintAttr() { auto attr = getProperties().l1_hint; getProperties().l1_hint = {}; return attr; } ::mlir::Attribute removeL2HintAttr() { auto attr = getProperties().l2_hint; getProperties().l2_hint = {}; return attr; } ::mlir::Attribute removeL3HintAttr() { auto attr = getProperties().l3_hint; getProperties().l3_hint = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value TensorDesc, ArrayRef offsets, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Value TensorDesc, ArrayRef offsets, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::ImplicitLocOpBuilder &builder, Value TensorDesc, ArrayRef offsets, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static PrefetchNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static PrefetchNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static PrefetchNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static PrefetchNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 5 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } xegpu::TensorDescType getTensorDescType() { return getTensorDesc().getType(); } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } SmallVector getMixedOffsets() { return getMixedValues(getConstOffsets(), getOffsets(), getContext()); } xegpu::DistributeLayoutAttr getDescLayoutAttr() { return dyn_cast_if_present(getTensorDescType().getLayout()); } ArrayRef getDataShape() { return getTensorDescType().getShape(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::PrefetchNdOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::PrefetchNdOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::PrefetchOp declarations //===----------------------------------------------------------------------===// namespace detail { class PrefetchOpGenericAdaptorBase { public: struct Properties { using l1_hintTy = ::mlir::xegpu::CachePolicyAttr; l1_hintTy l1_hint; auto getL1Hint() const { auto &propStorage = this->l1_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL1Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l1_hint = propValue; } using l2_hintTy = ::mlir::xegpu::CachePolicyAttr; l2_hintTy l2_hint; auto getL2Hint() const { auto &propStorage = this->l2_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL2Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l2_hint = propValue; } using l3_hintTy = ::mlir::xegpu::CachePolicyAttr; l3_hintTy l3_hint; auto getL3Hint() const { auto &propStorage = this->l3_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL3Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l3_hint = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } using offset_align_byteTy = ::mlir::IntegerAttr; offset_align_byteTy offset_align_byte; auto getOffsetAlignByte() const { auto &propStorage = this->offset_align_byte; return ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(propStorage); } void setOffsetAlignByte(const ::mlir::IntegerAttr &propValue) { this->offset_align_byte = propValue; } bool operator==(const Properties &rhs) const { return rhs.l1_hint == this->l1_hint && rhs.l2_hint == this->l2_hint && rhs.l3_hint == this->l3_hint && rhs.layout == this->layout && rhs.offset_align_byte == this->offset_align_byte && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: PrefetchOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.prefetch", odsAttrs.getContext()); } PrefetchOpGenericAdaptorBase(PrefetchOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::IntegerAttr getOffsetAlignByteAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(getProperties().offset_align_byte); return attr; } ::std::optional getOffsetAlignByte(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class PrefetchOpGenericAdaptor : public detail::PrefetchOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::PrefetchOpGenericAdaptorBase; static constexpr int odsIndex_source = 0; static constexpr int odsIndex_offsets = 1; public: PrefetchOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} PrefetchOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : PrefetchOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} PrefetchOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : PrefetchOpGenericAdaptor(values, attrs, Properties{}, {}) {} PrefetchOpGenericAdaptor(RangeT values, const PrefetchOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> PrefetchOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getSource() { return (*getODSOperands(0).begin()); } ValueT getOffsets() { return (*getODSOperands(1).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class PrefetchOpAdaptor : public PrefetchOpGenericAdaptor<::mlir::ValueRange> { public: using PrefetchOpGenericAdaptor::PrefetchOpGenericAdaptor; PrefetchOpAdaptor(PrefetchOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class PrefetchOp : public ::mlir::Op::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = PrefetchOpAdaptor; template using GenericAdaptor = PrefetchOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_source = 0; static constexpr int odsIndex_offsets = 1; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("l1_hint"), ::llvm::StringRef("l2_hint"), ::llvm::StringRef("l3_hint"), ::llvm::StringRef("layout"), ::llvm::StringRef("offset_align_byte")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getL1HintAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getL1HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getL2HintAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getL2HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getL3HintAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getL3HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(3); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 3); } ::mlir::StringAttr getOffsetAlignByteAttrName() { return getAttributeNameForIndex(4); } static ::mlir::StringAttr getOffsetAlignByteAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 4); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.prefetch"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getSource() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::Type> getOffsets() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(1).begin()); } ::mlir::OpOperand &getSourceMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getOffsetsMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::IntegerAttr getOffsetAlignByteAttr() { return ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(getProperties().offset_align_byte); } ::std::optional getOffsetAlignByte(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setL1HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l1_hint = attr; } void setL1Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL2HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l2_hint = attr; } void setL2Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL3HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l3_hint = attr; } void setL3Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setOffsetAlignByteAttr(::mlir::IntegerAttr attr) { getProperties().offset_align_byte = attr; } void setOffsetAlignByte(::std::optional attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeL1HintAttr() { auto attr = getProperties().l1_hint; getProperties().l1_hint = {}; return attr; } ::mlir::Attribute removeL2HintAttr() { auto attr = getProperties().l2_hint; getProperties().l2_hint = {}; return attr; } ::mlir::Attribute removeL3HintAttr() { auto attr = getProperties().l3_hint; getProperties().l3_hint = {}; return attr; } ::mlir::Attribute removeOffsetAlignByteAttr() { auto attr = getProperties().offset_align_byte; getProperties().offset_align_byte = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value source, ::mlir::Value offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::IntegerAttr offset_align_byte, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value source, ::mlir::Value offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::IntegerAttr offset_align_byte, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value source, ::mlir::Value offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::IntegerAttr offset_align_byte, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::Value offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::IntegerAttr offset_align_byte, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::Value offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::IntegerAttr offset_align_byte, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static PrefetchOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value source, ::mlir::Value offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::IntegerAttr offset_align_byte, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static PrefetchOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static PrefetchOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static PrefetchOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static PrefetchOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 5 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } Type getSourceType() { return getSource().getType(); } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::PrefetchOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::PrefetchOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::StoreMatrixOp declarations //===----------------------------------------------------------------------===// namespace detail { class StoreMatrixOpGenericAdaptorBase { public: struct Properties { using const_offsetsTy = ::mlir::DenseI64ArrayAttr; const_offsetsTy const_offsets; auto getConstOffsets() const { auto &propStorage = this->const_offsets; return ::llvm::cast<::mlir::DenseI64ArrayAttr>(propStorage); } void setConstOffsets(const ::mlir::DenseI64ArrayAttr &propValue) { this->const_offsets = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } using subgroup_block_ioTy = ::mlir::UnitAttr; subgroup_block_ioTy subgroup_block_io; auto getSubgroupBlockIo() const { auto &propStorage = this->subgroup_block_io; return ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(propStorage); } void setSubgroupBlockIo(const ::mlir::UnitAttr &propValue) { this->subgroup_block_io = propValue; } bool operator==(const Properties &rhs) const { return rhs.const_offsets == this->const_offsets && rhs.layout == this->layout && rhs.subgroup_block_io == this->subgroup_block_io && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: StoreMatrixOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.store_matrix", odsAttrs.getContext()); } StoreMatrixOpGenericAdaptorBase(StoreMatrixOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { auto attr = ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); return attr; } ::llvm::ArrayRef getConstOffsets(); ::mlir::UnitAttr getSubgroupBlockIoAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(getProperties().subgroup_block_io); return attr; } ::std::optional getSubgroupBlockIo(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class StoreMatrixOpGenericAdaptor : public detail::StoreMatrixOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::StoreMatrixOpGenericAdaptorBase; static constexpr int odsIndex_data = 0; static constexpr int odsIndex_mem_desc = 1; static constexpr int odsIndex_offsets = 2; public: StoreMatrixOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} StoreMatrixOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : StoreMatrixOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} StoreMatrixOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : StoreMatrixOpGenericAdaptor(values, attrs, Properties{}, {}) {} StoreMatrixOpGenericAdaptor(RangeT values, const StoreMatrixOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> StoreMatrixOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getData() { return (*getODSOperands(0).begin()); } ValueT getMemDesc() { return (*getODSOperands(1).begin()); } RangeT getOffsets() { return getODSOperands(2); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class StoreMatrixOpAdaptor : public StoreMatrixOpGenericAdaptor<::mlir::ValueRange> { public: using StoreMatrixOpGenericAdaptor::StoreMatrixOpGenericAdaptor; StoreMatrixOpAdaptor(StoreMatrixOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class StoreMatrixOp : public ::mlir::Op::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = StoreMatrixOpAdaptor; template using GenericAdaptor = StoreMatrixOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_data = 0; static constexpr int odsIndex_mem_desc = 1; static constexpr int odsIndex_offsets = 2; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("const_offsets"), ::llvm::StringRef("layout"), ::llvm::StringRef("subgroup_block_io")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getConstOffsetsAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getConstOffsetsAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getSubgroupBlockIoAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getSubgroupBlockIoAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.store_matrix"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getData() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::xegpu::MemDescType> getMemDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::MemDescType>>(*getODSOperands(1).begin()); } ::mlir::Operation::operand_range getOffsets() { return getODSOperands(2); } ::mlir::OpOperand &getDataMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getMemDescMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getOffsetsMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { return ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); } ::llvm::ArrayRef getConstOffsets(); ::mlir::UnitAttr getSubgroupBlockIoAttr() { return ::llvm::dyn_cast_or_null<::mlir::UnitAttr>(getProperties().subgroup_block_io); } ::std::optional getSubgroupBlockIo(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setConstOffsetsAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().const_offsets = attr; } void setConstOffsets(::llvm::ArrayRef attrValue); void setSubgroupBlockIoAttr(::mlir::UnitAttr attr) { getProperties().subgroup_block_io = attr; } void setSubgroupBlockIo(bool attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeSubgroupBlockIoAttr() { auto attr = getProperties().subgroup_block_io; getProperties().subgroup_block_io = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value data, TypedValue mem_desc, llvm::ArrayRef offsets, DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Value data, TypedValue mem_desc, llvm::ArrayRef offsets, DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, Value data, TypedValue mem_desc, llvm::ArrayRef offsets, DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value data, ::mlir::Value mem_desc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::UnitAttr subgroup_block_io, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static StoreMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static StoreMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static StoreMatrixOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static StoreMatrixOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 3 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: SmallVector getMixedOffsets() { return getMixedValues(getConstOffsets(), getOffsets(), getContext()); } ArrayRef getDataShape() { auto DataTy = getData().getType(); if (auto vecTy = llvm::dyn_cast(DataTy)) return vecTy.getShape(); return {}; } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::StoreMatrixOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::StoreMatrixOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::StoreNdOp declarations //===----------------------------------------------------------------------===// namespace detail { class StoreNdOpGenericAdaptorBase { public: struct Properties { using const_offsetsTy = ::mlir::DenseI64ArrayAttr; const_offsetsTy const_offsets; auto getConstOffsets() const { auto &propStorage = this->const_offsets; return ::llvm::cast<::mlir::DenseI64ArrayAttr>(propStorage); } void setConstOffsets(const ::mlir::DenseI64ArrayAttr &propValue) { this->const_offsets = propValue; } using l1_hintTy = ::mlir::xegpu::CachePolicyAttr; l1_hintTy l1_hint; auto getL1Hint() const { auto &propStorage = this->l1_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL1Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l1_hint = propValue; } using l2_hintTy = ::mlir::xegpu::CachePolicyAttr; l2_hintTy l2_hint; auto getL2Hint() const { auto &propStorage = this->l2_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL2Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l2_hint = propValue; } using l3_hintTy = ::mlir::xegpu::CachePolicyAttr; l3_hintTy l3_hint; auto getL3Hint() const { auto &propStorage = this->l3_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL3Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l3_hint = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } bool operator==(const Properties &rhs) const { return rhs.const_offsets == this->const_offsets && rhs.l1_hint == this->l1_hint && rhs.l2_hint == this->l2_hint && rhs.l3_hint == this->l3_hint && rhs.layout == this->layout && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: StoreNdOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.store_nd", odsAttrs.getContext()); } StoreNdOpGenericAdaptorBase(StoreNdOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize); const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { auto attr = ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); return attr; } ::llvm::ArrayRef getConstOffsets(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class StoreNdOpGenericAdaptor : public detail::StoreNdOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::StoreNdOpGenericAdaptorBase; static constexpr int odsIndex_value = 0; static constexpr int odsIndex_TensorDesc = 1; static constexpr int odsIndex_offsets = 2; public: StoreNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} StoreNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : StoreNdOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} StoreNdOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : StoreNdOpGenericAdaptor(values, attrs, Properties{}, {}) {} StoreNdOpGenericAdaptor(RangeT values, const StoreNdOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> StoreNdOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getValue() { return (*getODSOperands(0).begin()); } ValueT getTensorDesc() { return (*getODSOperands(1).begin()); } RangeT getOffsets() { return getODSOperands(2); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class StoreNdOpAdaptor : public StoreNdOpGenericAdaptor<::mlir::ValueRange> { public: using StoreNdOpGenericAdaptor::StoreNdOpGenericAdaptor; StoreNdOpAdaptor(StoreNdOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class StoreNdOp : public ::mlir::Op::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = StoreNdOpAdaptor; template using GenericAdaptor = StoreNdOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_value = 0; static constexpr int odsIndex_TensorDesc = 1; static constexpr int odsIndex_offsets = 2; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("const_offsets"), ::llvm::StringRef("l1_hint"), ::llvm::StringRef("l2_hint"), ::llvm::StringRef("l3_hint"), ::llvm::StringRef("layout")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getConstOffsetsAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getConstOffsetsAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getL1HintAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getL1HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getL2HintAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getL2HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } ::mlir::StringAttr getL3HintAttrName() { return getAttributeNameForIndex(3); } static ::mlir::StringAttr getL3HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 3); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(4); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 4); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.store_nd"); } std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getValue() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::xegpu::TensorDescType> getTensorDesc() { return ::llvm::cast<::mlir::TypedValue<::mlir::xegpu::TensorDescType>>(*getODSOperands(1).begin()); } ::mlir::Operation::operand_range getOffsets() { return getODSOperands(2); } ::mlir::OpOperand &getValueMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getTensorDescMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } ::mlir::MutableOperandRange getOffsetsMutable(); std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::DenseI64ArrayAttr getConstOffsetsAttr() { return ::llvm::cast<::mlir::DenseI64ArrayAttr>(getProperties().const_offsets); } ::llvm::ArrayRef getConstOffsets(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setConstOffsetsAttr(::mlir::DenseI64ArrayAttr attr) { getProperties().const_offsets = attr; } void setConstOffsets(::llvm::ArrayRef attrValue); void setL1HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l1_hint = attr; } void setL1Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL2HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l2_hint = attr; } void setL2Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL3HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l3_hint = attr; } void setL3Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeL1HintAttr() { auto attr = getProperties().l1_hint; getProperties().l1_hint = {}; return attr; } ::mlir::Attribute removeL2HintAttr() { auto attr = getProperties().l2_hint; getProperties().l2_hint = {}; return attr; } ::mlir::Attribute removeL3HintAttr() { auto attr = getProperties().l3_hint; getProperties().l3_hint = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value value, Value TensorDesc, ArrayRef offsets, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Value value, Value TensorDesc, ArrayRef offsets, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::ImplicitLocOpBuilder &builder, Value value, Value TensorDesc, ArrayRef offsets, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::mlir::DenseI64ArrayAttr const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value TensorDesc, ::mlir::ValueRange offsets, ::llvm::ArrayRef const_offsets, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static StoreNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static StoreNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static StoreNdOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static StoreNdOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 5 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } VectorType getValueType() { return llvm::dyn_cast(getValue().getType()); } xegpu::TensorDescType getTensorDescType() { return getTensorDesc().getType(); } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } SmallVector getMixedOffsets() { return getMixedValues(getConstOffsets(), getOffsets(), getContext()); } xegpu::DistributeLayoutAttr getDescLayoutAttr() { return dyn_cast_if_present(getTensorDescType().getLayout()); } ArrayRef getDataShape() { return getTensorDescType().getShape(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::StoreNdOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::StoreNdOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::StoreScatterOp declarations //===----------------------------------------------------------------------===// namespace detail { class StoreScatterOpGenericAdaptorBase { public: struct Properties { using chunk_sizeTy = ::mlir::IntegerAttr; chunk_sizeTy chunk_size; auto getChunkSize() const { auto &propStorage = this->chunk_size; return ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(propStorage); } void setChunkSize(const ::mlir::IntegerAttr &propValue) { this->chunk_size = propValue; } using l1_hintTy = ::mlir::xegpu::CachePolicyAttr; l1_hintTy l1_hint; auto getL1Hint() const { auto &propStorage = this->l1_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL1Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l1_hint = propValue; } using l2_hintTy = ::mlir::xegpu::CachePolicyAttr; l2_hintTy l2_hint; auto getL2Hint() const { auto &propStorage = this->l2_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL2Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l2_hint = propValue; } using l3_hintTy = ::mlir::xegpu::CachePolicyAttr; l3_hintTy l3_hint; auto getL3Hint() const { auto &propStorage = this->l3_hint; return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(propStorage); } void setL3Hint(const ::mlir::xegpu::CachePolicyAttr &propValue) { this->l3_hint = propValue; } using layoutTy = ::mlir::xegpu::DistributeLayoutAttr; layoutTy layout; auto getLayout() const { auto &propStorage = this->layout; return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(propStorage); } void setLayout(const ::mlir::xegpu::DistributeLayoutAttr &propValue) { this->layout = propValue; } bool operator==(const Properties &rhs) const { return rhs.chunk_size == this->chunk_size && rhs.l1_hint == this->l1_hint && rhs.l2_hint == this->l2_hint && rhs.l3_hint == this->l3_hint && rhs.layout == this->layout && true; } bool operator!=(const Properties &rhs) const { return !(*this == rhs); } }; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: StoreScatterOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.store", odsAttrs.getContext()); } StoreScatterOpGenericAdaptorBase(StoreScatterOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } ::mlir::IntegerAttr getChunkSizeAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(getProperties().chunk_size); return attr; } ::std::optional getChunkSize(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); return attr; } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { auto attr = ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); return attr; } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); }; } // namespace detail template class StoreScatterOpGenericAdaptor : public detail::StoreScatterOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::StoreScatterOpGenericAdaptorBase; static constexpr int odsIndex_value = 0; static constexpr int odsIndex_dest = 1; static constexpr int odsIndex_offsets = 2; static constexpr int odsIndex_mask = 3; public: StoreScatterOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} StoreScatterOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : StoreScatterOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} StoreScatterOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : StoreScatterOpGenericAdaptor(values, attrs, Properties{}, {}) {} StoreScatterOpGenericAdaptor(RangeT values, const StoreScatterOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> StoreScatterOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getValue() { return (*getODSOperands(0).begin()); } ValueT getDest() { return (*getODSOperands(1).begin()); } ValueT getOffsets() { return (*getODSOperands(2).begin()); } ValueT getMask() { return (*getODSOperands(3).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class StoreScatterOpAdaptor : public StoreScatterOpGenericAdaptor<::mlir::ValueRange> { public: using StoreScatterOpGenericAdaptor::StoreScatterOpGenericAdaptor; StoreScatterOpAdaptor(StoreScatterOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class StoreScatterOp : public ::mlir::Op::Impl, ::mlir::OpTrait::OpInvariants, ::mlir::BytecodeOpInterface::Trait, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::xegpu::AnchorLayoutInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = StoreScatterOpAdaptor; template using GenericAdaptor = StoreScatterOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_value = 0; static constexpr int odsIndex_dest = 1; static constexpr int odsIndex_offsets = 2; static constexpr int odsIndex_mask = 3; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { static ::llvm::StringRef attrNames[] = {::llvm::StringRef("chunk_size"), ::llvm::StringRef("l1_hint"), ::llvm::StringRef("l2_hint"), ::llvm::StringRef("l3_hint"), ::llvm::StringRef("layout")}; return ::llvm::ArrayRef(attrNames); } ::mlir::StringAttr getChunkSizeAttrName() { return getAttributeNameForIndex(0); } static ::mlir::StringAttr getChunkSizeAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 0); } ::mlir::StringAttr getL1HintAttrName() { return getAttributeNameForIndex(1); } static ::mlir::StringAttr getL1HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 1); } ::mlir::StringAttr getL2HintAttrName() { return getAttributeNameForIndex(2); } static ::mlir::StringAttr getL2HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 2); } ::mlir::StringAttr getL3HintAttrName() { return getAttributeNameForIndex(3); } static ::mlir::StringAttr getL3HintAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 3); } ::mlir::StringAttr getLayoutAttrName() { return getAttributeNameForIndex(4); } static ::mlir::StringAttr getLayoutAttrName(::mlir::OperationName name) { return getAttributeNameForIndex(name, 4); } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.store"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::Type> getValue() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(0).begin()); } ::mlir::TypedValue<::mlir::Type> getDest() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(1).begin()); } ::mlir::TypedValue<::mlir::Type> getOffsets() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(2).begin()); } ::mlir::TypedValue<::mlir::Type> getMask() { return ::llvm::cast<::mlir::TypedValue<::mlir::Type>>(*getODSOperands(3).begin()); } ::mlir::OpOperand &getValueMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getDestMutable() { auto range = getODSOperandIndexAndLength(1); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getOffsetsMutable() { auto range = getODSOperandIndexAndLength(2); return getOperation()->getOpOperand(range.first); } ::mlir::OpOperand &getMaskMutable() { auto range = getODSOperandIndexAndLength(3); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::llvm::LogicalResult readProperties(::mlir::DialectBytecodeReader &reader, ::mlir::OperationState &state); void writeProperties(::mlir::DialectBytecodeWriter &writer); ::mlir::IntegerAttr getChunkSizeAttr() { return ::llvm::dyn_cast_or_null<::mlir::IntegerAttr>(getProperties().chunk_size); } ::std::optional getChunkSize(); ::mlir::xegpu::CachePolicyAttr getL1HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l1_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL1Hint(); ::mlir::xegpu::CachePolicyAttr getL2HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l2_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL2Hint(); ::mlir::xegpu::CachePolicyAttr getL3HintAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::CachePolicyAttr>(getProperties().l3_hint); } ::std::optional<::mlir::xegpu::CachePolicy> getL3Hint(); ::mlir::xegpu::DistributeLayoutAttr getLayoutAttr() { return ::llvm::dyn_cast_or_null<::mlir::xegpu::DistributeLayoutAttr>(getProperties().layout); } ::std::optional<::mlir::xegpu::DistributeLayoutAttr> getLayout(); void setChunkSizeAttr(::mlir::IntegerAttr attr) { getProperties().chunk_size = attr; } void setChunkSize(::std::optional attrValue); void setL1HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l1_hint = attr; } void setL1Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL2HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l2_hint = attr; } void setL2Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setL3HintAttr(::mlir::xegpu::CachePolicyAttr attr) { getProperties().l3_hint = attr; } void setL3Hint(::std::optional<::mlir::xegpu::CachePolicy> attrValue); void setLayoutAttr(::mlir::xegpu::DistributeLayoutAttr attr) { getProperties().layout = attr; } ::mlir::Attribute removeChunkSizeAttr() { auto attr = getProperties().chunk_size; getProperties().chunk_size = {}; return attr; } ::mlir::Attribute removeL1HintAttr() { auto attr = getProperties().l1_hint; getProperties().l1_hint = {}; return attr; } ::mlir::Attribute removeL2HintAttr() { auto attr = getProperties().l2_hint; getProperties().l2_hint = {}; return attr; } ::mlir::Attribute removeL3HintAttr() { auto attr = getProperties().l3_hint; getProperties().l3_hint = {}; return attr; } ::mlir::Attribute removeLayoutAttr() { auto attr = getProperties().layout; getProperties().layout = {}; return attr; } static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value value, Value dest, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint); static StoreScatterOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Value value, Value dest, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint); static StoreScatterOp create(::mlir::ImplicitLocOpBuilder &builder, Value value, Value dest, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value value, Value dest, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static StoreScatterOp create(::mlir::OpBuilder &builder, ::mlir::Location location, Value value, Value dest, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static StoreScatterOp create(::mlir::ImplicitLocOpBuilder &builder, Value value, Value dest, ArrayRef offsets, Value mask, IntegerAttr chunk_size, xegpu::CachePolicyAttr l1_hint, xegpu::CachePolicyAttr l2_hint, xegpu::CachePolicyAttr l3_hint, xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value value, ::mlir::Value dest, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreScatterOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Value value, ::mlir::Value dest, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreScatterOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Value value, ::mlir::Value dest, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value dest, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreScatterOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value dest, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static StoreScatterOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value value, ::mlir::Value dest, ::mlir::Value offsets, ::mlir::Value mask, /*optional*/::mlir::IntegerAttr chunk_size, /*optional*/::mlir::xegpu::CachePolicyAttr l1_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l2_hint, /*optional*/::mlir::xegpu::CachePolicyAttr l3_hint, /*optional*/::mlir::xegpu::DistributeLayoutAttr layout); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static StoreScatterOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static StoreScatterOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static StoreScatterOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static StoreScatterOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); static ::llvm::LogicalResult setPropertiesFromParsedAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); private: ::mlir::StringAttr getAttributeNameForIndex(unsigned index) { return getAttributeNameForIndex((*this)->getName(), index); } static ::mlir::StringAttr getAttributeNameForIndex(::mlir::OperationName name, unsigned index) { assert(index < 5 && "invalid attribute index"); assert(name.getStringRef() == getOperationName() && "invalid operation name"); assert(name.isRegistered() && "Operation isn't registered, missing a " "dependent dialect loading?"); return name.getAttributeNames()[index]; } public: void printProperties(::mlir::MLIRContext *ctx, ::mlir::OpAsmPrinter &p, const Properties &prop, ::mlir::ArrayRef<::llvm::StringRef> elidedProps) { DictionaryAttr propAttr = dyn_cast_if_present(getPropertiesAsAttr(ctx, prop)); // filter out the elidedProps from propAttr, and get the resultAttr mlir::SmallVector filteredAttrs; if (propAttr) { for (auto namedAttr : propAttr.getValue()) { if (llvm::is_contained(elidedProps, namedAttr.getName().strref())) continue; filteredAttrs.push_back(namedAttr); } } if (!filteredAttrs.empty()) { p << "<" << DictionaryAttr::get(ctx, filteredAttrs) << ">"; } } static ::mlir::ParseResult parseProperties(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { if (mlir::succeeded(parser.parseOptionalLess())) { if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) return failure(); } return success(); } Type getDestType() { return getDest().getType(); } xegpu::DistributeLayoutAttr getAnchorLayout() { return getLayout().value_or(nullptr); } void setAnchorLayout(xegpu::DistributeLayoutAttr anchorLayout) { setLayoutAttr(anchorLayout); } mlir::Type getElementType() { auto type = getValue().getType(); return getElementTypeOrSelf(type); } VectorType getValueType() { return llvm::dyn_cast(getValue().getType()); } Type getMaskType() { return getMask().getType(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::StoreScatterOp) MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::detail::StoreScatterOpGenericAdaptorBase::Properties) namespace mlir::xegpu { //===----------------------------------------------------------------------===// // ::mlir::xegpu::TruncfOp declarations //===----------------------------------------------------------------------===// namespace detail { class TruncfOpGenericAdaptorBase { public: using Properties = ::mlir::EmptyProperties; protected: ::mlir::DictionaryAttr odsAttrs; ::std::optional<::mlir::OperationName> odsOpName; Properties properties; ::mlir::RegionRange odsRegions; public: TruncfOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), properties(properties), odsRegions(regions) { if (odsAttrs) odsOpName.emplace("xegpu.truncf", odsAttrs.getContext()); } TruncfOpGenericAdaptorBase(TruncfOp op); std::pair getODSOperandIndexAndLength(unsigned index, unsigned odsOperandsSize) { return {index, 1}; } const Properties &getProperties() { return properties; } ::mlir::DictionaryAttr getAttributes() { return odsAttrs; } }; } // namespace detail template class TruncfOpGenericAdaptor : public detail::TruncfOpGenericAdaptorBase { using ValueT = ::llvm::detail::ValueOfRange; using Base = detail::TruncfOpGenericAdaptorBase; static constexpr int odsIndex_source = 0; public: TruncfOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, const Properties &properties, ::mlir::RegionRange regions = {}) : Base(attrs, properties, regions), odsOperands(values) {} TruncfOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs, ::mlir::PropertyRef properties, ::mlir::RegionRange regions = {}) : TruncfOpGenericAdaptor(values, attrs, (properties ? *properties.as() : Properties{}), regions) {} TruncfOpGenericAdaptor(RangeT values, ::mlir::DictionaryAttr attrs = nullptr) : TruncfOpGenericAdaptor(values, attrs, Properties{}, {}) {} TruncfOpGenericAdaptor(RangeT values, const TruncfOpGenericAdaptorBase &base) : Base(base), odsOperands(values) {} template >> TruncfOpGenericAdaptor(RangeT values, LateInst op) : Base(op), odsOperands(values) {} std::pair getODSOperandIndexAndLength(unsigned index) { return Base::getODSOperandIndexAndLength(index, odsOperands.size()); } RangeT getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(odsOperands.begin(), valueRange.first), std::next(odsOperands.begin(), valueRange.first + valueRange.second)}; } ValueT getSource() { return (*getODSOperands(0).begin()); } RangeT getOperands() { return odsOperands; } private: RangeT odsOperands; }; class TruncfOpAdaptor : public TruncfOpGenericAdaptor<::mlir::ValueRange> { public: using TruncfOpGenericAdaptor::TruncfOpGenericAdaptor; TruncfOpAdaptor(TruncfOp op); ::llvm::LogicalResult verify(::mlir::Location loc); }; class TruncfOp : public ::mlir::Op::Impl, ::mlir::OpTrait::ZeroSuccessors, ::mlir::OpTrait::OneOperand, ::mlir::OpTrait::OpInvariants, ::mlir::ConditionallySpeculatable::Trait, ::mlir::OpTrait::AlwaysSpeculatableImplTrait, ::mlir::MemoryEffectOpInterface::Trait> { public: using Op::Op; using Op::print; using Adaptor = TruncfOpAdaptor; template using GenericAdaptor = TruncfOpGenericAdaptor; using FoldAdaptor = GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>; static constexpr int odsIndex_source = 0; using Properties = FoldAdaptor::Properties; static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() { return {}; } static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral("xegpu.truncf"); } std::pair getODSOperandIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::operand_range getODSOperands(unsigned index) { auto valueRange = getODSOperandIndexAndLength(index); return {std::next(getOperation()->operand_begin(), valueRange.first), std::next(getOperation()->operand_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getSource() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSOperands(0).begin()); } ::mlir::OpOperand &getSourceMutable() { auto range = getODSOperandIndexAndLength(0); return getOperation()->getOpOperand(range.first); } std::pair getODSResultIndexAndLength(unsigned index) { return {index, 1}; } ::mlir::Operation::result_range getODSResults(unsigned index) { auto valueRange = getODSResultIndexAndLength(index); return {std::next(getOperation()->result_begin(), valueRange.first), std::next(getOperation()->result_begin(), valueRange.first + valueRange.second)}; } ::mlir::TypedValue<::mlir::VectorType> getResult() { return ::llvm::cast<::mlir::TypedValue<::mlir::VectorType>>(*getODSResults(0).begin()); } static ::llvm::LogicalResult setPropertiesFromAttr(Properties &prop, ::mlir::Attribute attr, ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static ::mlir::Attribute getPropertiesAsAttr(::mlir::MLIRContext *ctx, const Properties &prop); static llvm::hash_code computePropertiesHash(const Properties &prop); static std::optional getInherentAttr(::mlir::MLIRContext *ctx, const Properties &prop, llvm::StringRef name); static void setInherentAttr(Properties &prop, llvm::StringRef name, mlir::Attribute value); static void populateInherentAttrs(::mlir::MLIRContext *ctx, const Properties &prop, ::mlir::NamedAttrList &attrs); static ::llvm::LogicalResult verifyInherentAttrs(::mlir::OperationName opName, ::mlir::NamedAttrList &attrs, llvm::function_ref<::mlir::InFlightDiagnostic()> emitError); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type result, ::mlir::Value source); static TruncfOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::Type result, ::mlir::Value source); static TruncfOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::Type result, ::mlir::Value source); static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value source); static TruncfOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::Value source); static TruncfOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::Value source); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static TruncfOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static TruncfOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {}); static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static TruncfOp create(::mlir::OpBuilder &builder, ::mlir::Location location, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); static TruncfOp create(::mlir::ImplicitLocOpBuilder &builder, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, const Properties &properties, ::llvm::ArrayRef<::mlir::NamedAttribute> discardableAttributes = {}); ::llvm::LogicalResult verifyInvariantsImpl(); ::llvm::LogicalResult verifyInvariants(); ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result); void print(::mlir::OpAsmPrinter &_odsPrinter); void getEffects(::llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<::mlir::MemoryEffects::Effect>> &effects); public: Type getSourceType() { return getSource().getType(); } Type getResultType() { return getResult().getType(); } }; } // namespace mlir::xegpu MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::xegpu::TruncfOp) #endif // GET_OP_CLASSES