diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td index dabbfd4c79de..2b977fc41855 100644 --- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td +++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td @@ -132,6 +132,27 @@ def XeGPU_FenceScopeAttr: let assemblyFormat = "$value"; } +//===----------------------------------------------------------------------===// +// XeGPU Lane Shuffle Enums. +//===----------------------------------------------------------------------===// +def XeGPU_LaneShufflePack: I32EnumAttrCase<"Pack", 0, "pack">; +def XeGPU_LaneShuffleUnpack: I32EnumAttrCase<"Unpack", 1, "unpack">; +def XeGPU_LaneShuffleMode: I32EnumAttr<"LaneShuffleMode", + "The direction of the lane shuffle performed by xegpu.lane_shuffle.", + [XeGPU_LaneShufflePack, XeGPU_LaneShuffleUnpack]> { + let genSpecializedAttr = 0; + let cppNamespace = "::mlir::xegpu"; +} + +def XeGPU_LaneShuffleModeAttr: + EnumAttr { + let summary = [{Describes the direction of a lane shuffle. + "pack" gathers a lane's elements into consecutive logical + positions. "unpack" scatters them back out, strided by the + subgroup size.}]; + let assemblyFormat = "$value"; +} + def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> { let cppNamespace = "::mlir::xegpu"; let description = [{ diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td index 65c046bb37ac..49b98922cee4 100644 --- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td +++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td @@ -1500,6 +1500,78 @@ def XeGPU_TruncfOp let hasVerifier = 1; } +def XeGPU_LaneShuffleOp + : XeGPU_Op<"lane_shuffle", [Pure, AllTypesMatch<["source", "result"]>]> { + let summary = "Re-distributes a subgroup's fragments across its lanes."; + + let description = [{ + `xegpu.lane_shuffle` is a lane-level operation. The `source` operand is the + fragment (a 1D vector) held by one lane, and the `result` is the fragment + that lane holds afterwards. It re-distributes the elements across the lanes + of the subgroup without changing their type, so `source` and `result` have + the same type. + + Let `S` be the subgroup size and `N` the number of elements per lane, so the + subgroup collectively holds `S * N` elements. Number these elements by their + logical position `0 .. S*N-1`. The `mode` selects which of the two + orientations of the `S x N` (lane, element) index grid the operation maps + between: + + * `pack`: element `j` of lane `i` holds logical position `j * S + i` before + the operation and `i * N + j` after it. A lane's elements end up at + consecutive logical positions. + * `unpack`: the reverse mapping. A lane's elements end up strided by `S`. + + The two modes are exact inverses, so shuffling one way and then the other + yields the original fragments. + + For example, consider `vector<2xi16>` with a subgroup size of 8, denoting a + lane's fragment as ``. In `pack` mode: + + ``` + lane: 0 1 2 3 4 5 6 7 + source: <0,8> <1,9> <2,10> <3,11> <4,12> <5,13> <6,14> <7,15> + result: <0,1> <2,3> <4,5> <6,7> <8,9> <10,11> <12,13> <14,15> + ``` + + In general each lane's result elements come from `N` different lanes' + `source` fragments -- that is the cross-lane shuffle. + + This operation implements the `xegpu.convert_layout` semantics at the lane + level, for the case where the `lane_layout` is the same but the `lane_data` + is different. + + The operation assumes a full subgroup (the subgroup size equals the maximum + subgroup size) and that all lanes execute the same dynamic instance of the + operation; otherwise the behavior is undefined. + + Arguments: + - `source`: The 1D vector fragment held by one lane. + - `mode`: The direction of the shuffle, either `pack` or `unpack`. + }]; + + let arguments = (ins XeGPU_LaneShuffleType:$source, + XeGPU_LaneShuffleModeAttr:$mode); + let results = (outs XeGPU_LaneShuffleType:$result); + + let extraClassDeclaration = [{ + VectorType getSourceType() { + return llvm::cast(getSource().getType()); + } + + VectorType getResultType() { + return llvm::cast(getResult().getType()); + } + }]; + + let assemblyFormat = [{ + $source $mode attr-dict `:` type($source) + }]; + + let hasVerifier = 1; + let hasFolder = 1; +} + def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments, AnchorLayoutInterface]> { let summary = "It performs scaled mma computation"; diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td index 78c29ae9e7f4..ec0e083aa207 100644 --- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td +++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td @@ -31,6 +31,9 @@ def XeGPU_VectorOrScalarType : AnyTypeOf<[VectorOfRankAndType<[1,2,3,4,5,6,7,8], [XeGPU_ScalarType, Index]>, XeGPU_ScalarType]>; def XeGPU_GatherScatterBaseAddrType : AnyTypeOf<[MemRefRankOf<[XeGPU_ScalarType], [1]>, XeGPU_PointerType]>; +// Operand/result type of xegpu.lane_shuffle: a 1D vector of numerical type. +def XeGPU_LaneShuffleType + : FixedVectorOfRankAndType<[1], [XeGPU_ScalarType]>; // common base class for types in XeGPU dialect class XeGPUTypeDef traits = [], diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp index 060da6230286..1bd5951c9f7f 100644 --- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp +++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp @@ -946,6 +946,29 @@ LogicalResult TruncfOp::verify() { return success(); } +//===----------------------------------------------------------------------===// +// XeGPU_LaneShuffleOp +//===----------------------------------------------------------------------===// + +LogicalResult LaneShuffleOp::verify() { + // With a single element per lane there is nothing to re-distribute, so the + // operation would be a no-op. + if (getSourceType().getNumElements() < 2) + return emitOpError("requires a source vector with at least 2 elements."); + + return success(); +} + +OpFoldResult LaneShuffleOp::fold(FoldAdaptor adaptor) { + // The two modes are exact inverses, so a pack feeding an unpack (or vice + // versa) restores the original fragments. + auto producer = getSource().getDefiningOp(); + if (producer && producer.getMode() != getMode()) + return producer.getSource(); + + return {}; +} + //===----------------------------------------------------------------------===// // XeGPU_DpasMxOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/XeGPU/canonicalize.mlir b/mlir/test/Dialect/XeGPU/canonicalize.mlir new file mode 100644 index 000000000000..755d27ec4287 --- /dev/null +++ b/mlir/test/Dialect/XeGPU/canonicalize.mlir @@ -0,0 +1,36 @@ +// RUN: mlir-opt --canonicalize --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @fold_lane_shuffle_pack_unpack +// CHECK-SAME: %[[ARG0:.+]]: vector<2xi16> +// CHECK-NOT: xegpu.lane_shuffle +// CHECK: return %[[ARG0]] +func.func @fold_lane_shuffle_pack_unpack(%a: vector<2xi16>) -> vector<2xi16> { + %0 = xegpu.lane_shuffle %a pack : vector<2xi16> + %1 = xegpu.lane_shuffle %0 unpack : vector<2xi16> + return %1 : vector<2xi16> +} + +// ----- + +// CHECK-LABEL: func.func @fold_lane_shuffle_unpack_pack +// CHECK-SAME: %[[ARG0:.+]]: vector<4xf8E5M2> +// CHECK-NOT: xegpu.lane_shuffle +// CHECK: return %[[ARG0]] +func.func @fold_lane_shuffle_unpack_pack(%a: vector<4xf8E5M2>) -> vector<4xf8E5M2> { + %0 = xegpu.lane_shuffle %a unpack : vector<4xf8E5M2> + %1 = xegpu.lane_shuffle %0 pack : vector<4xf8E5M2> + return %1 : vector<4xf8E5M2> +} + +// ----- + +// Two shuffles in the same direction are not inverses and must not fold. + +// CHECK-LABEL: func.func @no_fold_lane_shuffle_pack_pack +// CHECK: xegpu.lane_shuffle {{.*}} pack +// CHECK: xegpu.lane_shuffle {{.*}} pack +func.func @no_fold_lane_shuffle_pack_pack(%a: vector<2xi16>) -> vector<2xi16> { + %0 = xegpu.lane_shuffle %a pack : vector<2xi16> + %1 = xegpu.lane_shuffle %0 pack : vector<2xi16> + return %1 : vector<2xi16> +} diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir index ac87045de701..dc68f5136b4a 100644 --- a/mlir/test/Dialect/XeGPU/invalid.mlir +++ b/mlir/test/Dialect/XeGPU/invalid.mlir @@ -693,6 +693,13 @@ func.func @truncf_invalid_result_size(%a: vector<8x16xf16>) { return } +// ----- +func.func @lane_shuffle_single_element(%a: vector<1xi32>) { + // expected-error@+1 {{op requires a source vector with at least 2 elements}} + %1 = xegpu.lane_shuffle %a pack : vector<1xi32> + return +} + // ----- func.func @dpas_mx_acc_result_type_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xbf16>) { // expected-error@+1 {{Accumulator type must match result type.}} diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir index 5591eea00a2e..f733491bfb7e 100644 --- a/mlir/test/Dialect/XeGPU/ops.mlir +++ b/mlir/test/Dialect/XeGPU/ops.mlir @@ -697,6 +697,48 @@ gpu.func @truncf(%a: vector<8x16xf16>) { gpu.return } +// CHECK-LABEL: gpu.func @lane_shuffle_pack_f4 +gpu.func @lane_shuffle_pack_f4(%a: vector<4xf4E2M1FN>) { + // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<4xf4E2M1FN> + %1 = xegpu.lane_shuffle %a pack : vector<4xf4E2M1FN> + gpu.return +} + +// CHECK-LABEL: gpu.func @lane_shuffle_unpack_f4 +gpu.func @lane_shuffle_unpack_f4(%a: vector<4xf4E2M1FN>) { + // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<4xf4E2M1FN> + %1 = xegpu.lane_shuffle %a unpack : vector<4xf4E2M1FN> + gpu.return +} + +// CHECK-LABEL: gpu.func @lane_shuffle_pack_f8 +gpu.func @lane_shuffle_pack_f8(%a: vector<4xf8E5M2>) { + // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<4xf8E5M2> + %1 = xegpu.lane_shuffle %a pack : vector<4xf8E5M2> + gpu.return +} + +// CHECK-LABEL: gpu.func @lane_shuffle_unpack_f8 +gpu.func @lane_shuffle_unpack_f8(%a: vector<4xf8E5M2>) { + // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<4xf8E5M2> + %1 = xegpu.lane_shuffle %a unpack : vector<4xf8E5M2> + gpu.return +} + +// CHECK-LABEL: gpu.func @lane_shuffle_pack_f16 +gpu.func @lane_shuffle_pack_f16(%a: vector<2xf16>) { + // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} pack : vector<2xf16> + %1 = xegpu.lane_shuffle %a pack : vector<2xf16> + gpu.return +} + +// CHECK-LABEL: gpu.func @lane_shuffle_unpack_f16 +gpu.func @lane_shuffle_unpack_f16(%a: vector<2xf16>) { + // CHECK: %{{.+}} = xegpu.lane_shuffle %{{.+}} unpack : vector<2xf16> + %1 = xegpu.lane_shuffle %a unpack : vector<2xf16> + gpu.return +} + // CHECK-LABEL: gpu.func @dpas_mx gpu.func @dpas_mx(%a : vector<8x32xf8E5M2>, %b: vector<32x16xf8E5M2>, %acc: vector<8x16xbf16>, %a_scale: vector<8x1xf8E8M0FNU>, %b_scale: vector<1x16xf8E8M0FNU>) { // CHECK: %{{.+}} = xegpu.dpas_mx %{{.+}}, %{{.+}}, %{{.+}} scale_a = %{{.+}} scale_b = %{{.+}} : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>