[mlir][spirv] Define sp.VectorShuffle

[antiagainst]

Summary

This patch adds basic op definition, parser/printer, and verifier.

[ThomasRaoux]

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td

295–296

Do we need a check to make sure the sources and result have the same component type? I don't see it.

[ergawy]

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td

332

Should it be SPV_Vector instead?

[antiagainst]

Address comments

[antiagainst]

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td

295–296

Good catch! I added a rule via AllElementTypesMatch. Generally for rules enforced by ODS directly, we can omit dedicated tests given they are already enforced by the framework. :)

332

Good catch! Done.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td

 def SPV_OC_OpLoad                      : I32EnumAttrCase<"OpLoad", 61>;
 def SPV_OC_OpStore                     : I32EnumAttrCase<"OpStore", 62>;
 def SPV_OC_OpCopyMemory                : I32EnumAttrCase<"OpCopyMemory", 63>;
 def SPV_OC_OpAccessChain               : I32EnumAttrCase<"OpAccessChain", 65>;
 def SPV_OC_OpDecorate                  : I32EnumAttrCase<"OpDecorate", 71>;
 def SPV_OC_OpMemberDecorate            : I32EnumAttrCase<"OpMemberDecorate", 72>;
 def SPV_OC_OpVectorExtractDynamic      : I32EnumAttrCase<"OpVectorExtractDynamic", 77>;
 def SPV_OC_OpVectorInsertDynamic       : I32EnumAttrCase<"OpVectorInsertDynamic", 78>;
+def SPV_OC_OpVectorShuffle             : I32EnumAttrCase<"OpVectorShuffle", 79>;
 def SPV_OC_OpCompositeConstruct        : I32EnumAttrCase<"OpCompositeConstruct", 80>;
 def SPV_OC_OpCompositeExtract          : I32EnumAttrCase<"OpCompositeExtract", 81>;
 def SPV_OC_OpCompositeInsert           : I32EnumAttrCase<"OpCompositeInsert", 82>;
 def SPV_OC_OpTranspose                 : I32EnumAttrCase<"OpTranspose", 84>;
 def SPV_OC_OpConvertFToU               : I32EnumAttrCase<"OpConvertFToU", 109>;
 def SPV_OC_OpConvertFToS               : I32EnumAttrCase<"OpConvertFToS", 110>;
 def SPV_OC_OpConvertSToF               : I32EnumAttrCase<"OpConvertSToF", 111>;
 def SPV_OC_OpConvertUToF               : I32EnumAttrCase<"OpConvertUToF", 112>;
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpConstantTrue, SPV_OC_OpConstantFalse, SPV_OC_OpConstant,
       SPV_OC_OpConstantComposite, SPV_OC_OpConstantNull, SPV_OC_OpSpecConstantTrue,
       SPV_OC_OpSpecConstantFalse, SPV_OC_OpSpecConstant,
       SPV_OC_OpSpecConstantComposite, SPV_OC_OpSpecConstantOp, SPV_OC_OpFunction,
       SPV_OC_OpFunctionParameter, SPV_OC_OpFunctionEnd, SPV_OC_OpFunctionCall,
       SPV_OC_OpVariable, SPV_OC_OpLoad, SPV_OC_OpStore, SPV_OC_OpCopyMemory,
       SPV_OC_OpAccessChain, SPV_OC_OpDecorate, SPV_OC_OpMemberDecorate,
       SPV_OC_OpVectorExtractDynamic, SPV_OC_OpVectorInsertDynamic,
-      SPV_OC_OpCompositeConstruct, SPV_OC_OpCompositeExtract,
+      SPV_OC_OpVectorShuffle, SPV_OC_OpCompositeConstruct, SPV_OC_OpCompositeExtract,
       SPV_OC_OpCompositeInsert, SPV_OC_OpTranspose, SPV_OC_OpConvertFToU,
       SPV_OC_OpConvertFToS, SPV_OC_OpConvertSToF, SPV_OC_OpConvertUToF,
       SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert, SPV_OC_OpBitcast,
       SPV_OC_OpSNegate, SPV_OC_OpFNegate, SPV_OC_OpIAdd, SPV_OC_OpFAdd,
       SPV_OC_OpISub, SPV_OC_OpFSub, SPV_OC_OpIMul, SPV_OC_OpFMul, SPV_OC_OpUDiv,
       SPV_OC_OpSDiv, SPV_OC_OpFDiv, SPV_OC_OpUMod, SPV_OC_OpSRem, SPV_OC_OpSMod,
       SPV_OC_OpFRem, SPV_OC_OpFMod, SPV_OC_OpMatrixTimesScalar,
       SPV_OC_OpMatrixTimesMatrix, SPV_OC_OpIsNan, SPV_OC_OpIsInf, SPV_OC_OpOrdered,

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td

def SPV_CompositeInsertOp : SPV_Op<"CompositeInsert",
   );
 
   let results = (outs
     SPV_Composite:$result
   );
 
   let builders = [
     OpBuilderDAG<(ins "Value":$object, "Value":$composite,
                       "ArrayRef<int32_t>":$indices)>
   ];
 }
 
 // -----
 
-def SPV_VectorExtractDynamicOp : SPV_Op<"VectorExtractDynamic",
-            [NoSideEffect,
-            TypesMatchWith<"type of 'result' matches element type of 'vector'",
-                     "vector", "result",
-                     "$_self.cast<mlir::VectorType>().getElementType()">]> {
+def SPV_VectorExtractDynamicOp : SPV_Op<"VectorExtractDynamic", [
+    NoSideEffect,
   let summary = [{
     Extract a single, dynamically selected, component of a vector.
   }];
 
   let description = [{
     Result Type must be a scalar type.
 
     Vector must have a type OpTypeVector whose Component Type is Result
     Type.
 
     Index must be a scalar integer. It is interpreted as a 0-based index of
     which component of Vector to extract.
 
     Behavior is undefined if Index's value is less than zero or greater than
     or equal to the number of components in Vector.
 
     <!-- End of AutoGen section -->
 
-    ```
-  scalar-type ::= integer-type | float-type | boolean-type 
-    vector-extract-dynamic-op ::= `spv.VectorExtractDynamic ` ssa-use `[` ssa-use `]`
-                                    `:` `vector<` integer-literal `x` scalar-type `>` `,`
-                                    integer-type
-    ```mlir
-
     #### Example:
 
     ```
     %2 = spv.VectorExtractDynamic %0[%1] : vector<8xf32>, i32
     ```
   }];
 
   let arguments = (ins
def SPV_VectorExtractDynamicOp : SPV_Op<"VectorExtractDynamic", [
   
   let assemblyFormat = [{
     $vector `[` $index `]` attr-dict `:` type($vector) `,` type($index) 
   }];
 }
 
 // -----
 
-def SPV_VectorInsertDynamicOp : SPV_Op<"VectorInsertDynamic",
-        [NoSideEffect,
-        TypesMatchWith<"type of 'component' matches element type of 'vector'",
-                "vector", "component",
-                "$_self.cast<mlir::VectorType>().getElementType()">,
-                AllTypesMatch<["vector", "result"]>]> {
+def SPV_VectorInsertDynamicOp : SPV_Op<"VectorInsertDynamic", [
+    NoSideEffect,
+    TypesMatchWith<
+      "type of 'component' matches element type of 'vector'",
   let summary = [{
     Make a copy of a vector, with a single, variably selected, component
     modified.
   }];
 
   let description = [{
     Result Type must be an OpTypeVector.
 
def SPV_VectorInsertDynamicOp : SPV_Op<"VectorInsertDynamic", [
   
   let assemblyFormat = [{
     $component `,` $vector `[` $index `]` attr-dict `:` type($vector) `,` type($index) 
   }];
 }
 
 // -----
 
+def SPV_VectorShuffleOp : SPV_Op<"VectorShuffle", [
+    NoSideEffect, AllElementTypesMatch<["vector1", "vector2", "result"]>]> {
+  let summary = [{
+    Select arbitrary components from two vectors to make a new vector.
+  }];
+
+  let description = [{
+    Result Type must be an OpTypeVector. The number of components in Result
+    Type must be the same as the number of Component operands.
+
+    Vector 1 and Vector 2 must both have vector types, with the same
+    Component Type as Result Type. They do not have to have the same number
+    of components as Result Type or with each other. They are logically
+    concatenated, forming a single vector with Vector 1’s components
+    appearing before Vector 2’s. The components of this logical vector are
+    logically numbered with a single consecutive set of numbers from 0 to N
+    - 1, where N is the total number of components.
+
+    Components are these logical numbers (see above), selecting which of the
+    logically numbered components form the result. Each component is an
+    unsigned 32-bit integer.  They can select the components in any order
+    and can repeat components. The first component of the result is selected
+    by the first Component operand,  the second component of the result is
+    selected by the second Component operand, etc. A Component literal may
+    also be FFFFFFFF, which means the corresponding result component has no
+    source and is undefined. All Component literals must either be FFFFFFFF
+    or in [0, N - 1] (inclusive).
+
+    Note: A vector “swizzle” can be done by using the vector for both Vector
+    operands, or using an OpUndef for one of the Vector operands.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    %0 = spv.VectorShuffle [1: i32, 3: i32, 5: i32]
+                           %vector1: vector<4xf32>, %vector2: vector<2xf32>
+                        -> vector<3xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_Vector:$vector1,
+    SPV_Vector:$vector2,
+    I32ArrayAttr:$components
+  );
+
+  let results = (outs
+    SPV_Vector:$result
+  );
+
+  let assemblyFormat = [{
+    attr-dict $components $vector1 `:` type($vector1) `,`
+                          $vector2 `:` type($vector2) `->` type($result)
+  }];
+}
+
+// -----
+
 #endif // MLIR_DIALECT_SPIRV_IR_COMPOSITE_OPS

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/bit.h"
 
 using namespace mlir;
 
if (op->getAttr(attr))
       return varOp.emitOpError("cannot have '")
              << attr << "' attribute (only allowed in spv.globalVariable)";
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// spv.VectorShuffle
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(spirv::VectorShuffleOp shuffleOp) {
+  VectorType resultType = shuffleOp.getType().cast<VectorType>();
+
+  size_t numResultElements = resultType.getNumElements();
+  if (numResultElements != shuffleOp.components().size())
+    return shuffleOp.emitOpError("result type element count (")
+           << numResultElements
+           << ") mismatch with the number of component selectors ("
+           << shuffleOp.components().size() << ")";
+
+  size_t totalSrcElements =
+      shuffleOp.vector1().getType().cast<VectorType>().getNumElements() +
+      shuffleOp.vector2().getType().cast<VectorType>().getNumElements();
+
+  for (const auto &selector :
+       shuffleOp.components().getAsValueRange<IntegerAttr>()) {
+    uint32_t index = selector.getZExtValue();
+    if (index >= totalSrcElements &&
+        index != std::numeric_limits<uint32_t>().max())
+      return shuffleOp.emitOpError("component selector ")
+             << index << " out of range: expected to be in [0, "
+             << totalSrcElements << ") or 0xffffffff";
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // spv.CooperativeMatrixLoadNV
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseCooperativeMatrixLoadNVOp(OpAsmParser &parser,
                                                   OperationState &state) {
   SmallVector<OpAsmParser::OperandType, 3> operandInfo;
   Type strideType = parser.getBuilder().getIntegerType(32);
   Type columnMajorType = parser.getBuilder().getIntegerType(1);

mlir/test/Dialect/SPIRV/IR/composite-ops.mlir

 // spv.VectorInsertDynamic
 //===----------------------------------------------------------------------===//
 
 func @vector_dynamic_insert(%val: f32, %vec: vector<4xf32>, %id : i32) -> vector<4xf32> {
   // CHECK: spv.VectorInsertDynamic %{{.*}}, %{{.*}}[%{{.*}}] : vector<4xf32>, i32
   %0 = spv.VectorInsertDynamic %val, %vec[%id] : vector<4xf32>, i32
   return %0 : vector<4xf32>
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.VectorShuffle
+//===----------------------------------------------------------------------===//
+
+func @vector_shuffle(%vector1: vector<4xf32>, %vector2: vector<2xf32>) -> vector<3xf32> {
+  // CHECK: %{{.+}} = spv.VectorShuffle [1 : i32, 3 : i32, -1 : i32] %{{.+}} : vector<4xf32>, %arg1 : vector<2xf32> -> vector<3xf32>
+  %0 = spv.VectorShuffle [1: i32, 3: i32, 0xffffffff: i32] %vector1: vector<4xf32>, %vector2: vector<2xf32> -> vector<3xf32>
+  return %0: vector<3xf32>
+}
+
+// -----
+
+func @vector_shuffle_extra_selector(%vector1: vector<4xf32>, %vector2: vector<2xf32>) -> vector<3xf32> {
+  // expected-error @+1 {{result type element count (3) mismatch with the number of component selectors (4)}}
+  %0 = spv.VectorShuffle [1: i32, 3: i32, 5: i32, 2: i32] %vector1: vector<4xf32>, %vector2: vector<2xf32> -> vector<3xf32>
+  return %0: vector<3xf32>
+}
+
+// -----
+
+func @vector_shuffle_extra_selector(%vector1: vector<4xf32>, %vector2: vector<2xf32>) -> vector<3xf32> {
+  // expected-error @+1 {{component selector 7 out of range: expected to be in [0, 6) or 0xffffffff}}
+  %0 = spv.VectorShuffle [1: i32, 7: i32, 5: i32] %vector1: vector<4xf32>, %vector2: vector<2xf32> -> vector<3xf32>
+  return %0: vector<3xf32>
+}

mlir/test/Target/SPIRV/composite-op.mlir

spv.func @vector_dynamic_extract(%vec: vector<4xf32>, %id : i32) -> f32 "None" {
     %0 = spv.VectorExtractDynamic %vec[%id] : vector<4xf32>, i32
     spv.ReturnValue %0: f32
   }
   spv.func @vector_dynamic_insert(%val: f32, %vec: vector<4xf32>, %id : i32) -> vector<4xf32> "None" {
     // CHECK: spv.VectorInsertDynamic %{{.*}}, %{{.*}}[%{{.*}}] : vector<4xf32>, i32
     %0 = spv.VectorInsertDynamic %val, %vec[%id] : vector<4xf32>, i32
     spv.ReturnValue %0: vector<4xf32>
   }
+  spv.func @vector_shuffle(%vector1: vector<4xf32>, %vector2: vector<2xf32>) -> vector<3xf32> "None" {
+    // CHECK: %{{.+}} = spv.VectorShuffle [1 : i32, 3 : i32, -1 : i32] %{{.+}} : vector<4xf32>, %arg1 : vector<2xf32> -> vector<3xf32>
+    %0 = spv.VectorShuffle [1: i32, 3: i32, 0xffffffff: i32] %vector1: vector<4xf32>, %vector2: vector<2xf32> -> vector<3xf32>
+    spv.ReturnValue %0: vector<3xf32>
+  }
 }