[mlir][spirv] Add support for OpPtrCastToGeneric

[alexbatashev]

Summary

OpPtrCastToGeneric converts pointer's storage class to Generic.

[antiagainst]

Thanks for adding this!

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

515

This is an op in the core spec so it should not be placed in SPIRVOCLOps.td. Please define it in SPIRVCastOps.td.

543

Can use SPV_AnyPtr here.

547

Can use SPV_AnyPtr here.

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

4431

You can define

let assemblyFormat = "$pointer attr-dict `:` type($pointer) `to` type($result)";

in ODS to autogen the parser and printer.

4464

This will be autogen'ed with SPV_AnyPtr in the ODS.

4468

I think you want to check Workgroup, CrossWorkgroup, or Function here?

4472

This won't be needed with SPV_AnyPtr in the ODS.

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

194

This should be // CHECK: to actually checking it.

196

Please add negative tests triggering validation errors.

Also please add a test for (de)serialization in https://github.com/llvm/llvm-project/blob/main/mlir/test/Target/SPIRV/cast-ops.mlir.

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

 def SPV_OC_OpImageQuerySize            : I32EnumAttrCase<"OpImageQuerySize", 104>;
 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>;
 def SPV_OC_OpUConvert                  : I32EnumAttrCase<"OpUConvert", 113>;
 def SPV_OC_OpSConvert                  : I32EnumAttrCase<"OpSConvert", 114>;
 def SPV_OC_OpFConvert                  : I32EnumAttrCase<"OpFConvert", 115>;
+def SPV_OC_OpPtrCastToGeneric          : I32EnumAttrCase<"OpPtrCastToGeneric", 121>;
 def SPV_OC_OpBitcast                   : I32EnumAttrCase<"OpBitcast", 124>;
 def SPV_OC_OpSNegate                   : I32EnumAttrCase<"OpSNegate", 126>;
 def SPV_OC_OpFNegate                   : I32EnumAttrCase<"OpFNegate", 127>;
 def SPV_OC_OpIAdd                      : I32EnumAttrCase<"OpIAdd", 128>;
 def SPV_OC_OpFAdd                      : I32EnumAttrCase<"OpFAdd", 129>;
 def SPV_OC_OpISub                      : I32EnumAttrCase<"OpISub", 130>;
 def SPV_OC_OpFSub                      : I32EnumAttrCase<"OpFSub", 131>;
 def SPV_OC_OpIMul                      : I32EnumAttrCase<"OpIMul", 132>;
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpFunctionCall, SPV_OC_OpVariable, SPV_OC_OpLoad, SPV_OC_OpStore,
       SPV_OC_OpCopyMemory, SPV_OC_OpAccessChain, SPV_OC_OpPtrAccessChain,
       SPV_OC_OpInBoundsPtrAccessChain, SPV_OC_OpDecorate, SPV_OC_OpMemberDecorate,
       SPV_OC_OpVectorExtractDynamic, SPV_OC_OpVectorInsertDynamic,
       SPV_OC_OpVectorShuffle, SPV_OC_OpCompositeConstruct, SPV_OC_OpCompositeExtract,
       SPV_OC_OpCompositeInsert, SPV_OC_OpTranspose, SPV_OC_OpImageDrefGather,
       SPV_OC_OpImage, SPV_OC_OpImageQuerySize, 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,
-      SPV_OC_OpUnordered, SPV_OC_OpLogicalEqual, SPV_OC_OpLogicalNotEqual,
-      SPV_OC_OpLogicalOr, SPV_OC_OpLogicalAnd, SPV_OC_OpLogicalNot, SPV_OC_OpSelect,
-      SPV_OC_OpIEqual, SPV_OC_OpINotEqual, SPV_OC_OpUGreaterThan,
-      SPV_OC_OpSGreaterThan, SPV_OC_OpUGreaterThanEqual, SPV_OC_OpSGreaterThanEqual,
-      SPV_OC_OpULessThan, SPV_OC_OpSLessThan, SPV_OC_OpULessThanEqual,
-      SPV_OC_OpSLessThanEqual, SPV_OC_OpFOrdEqual, SPV_OC_OpFUnordEqual,
-      SPV_OC_OpFOrdNotEqual, SPV_OC_OpFUnordNotEqual, SPV_OC_OpFOrdLessThan,
-      SPV_OC_OpFUnordLessThan, SPV_OC_OpFOrdGreaterThan, SPV_OC_OpFUnordGreaterThan,
-      SPV_OC_OpFOrdLessThanEqual, SPV_OC_OpFUnordLessThanEqual,
-      SPV_OC_OpFOrdGreaterThanEqual, SPV_OC_OpFUnordGreaterThanEqual,
-      SPV_OC_OpShiftRightLogical, SPV_OC_OpShiftRightArithmetic,
-      SPV_OC_OpShiftLeftLogical, SPV_OC_OpBitwiseOr, SPV_OC_OpBitwiseXor,
-      SPV_OC_OpBitwiseAnd, SPV_OC_OpNot, SPV_OC_OpBitFieldInsert,
-      SPV_OC_OpBitFieldSExtract, SPV_OC_OpBitFieldUExtract, SPV_OC_OpBitReverse,
-      SPV_OC_OpBitCount, SPV_OC_OpControlBarrier, SPV_OC_OpMemoryBarrier,
-      SPV_OC_OpAtomicExchange, SPV_OC_OpAtomicCompareExchange,
-      SPV_OC_OpAtomicCompareExchangeWeak, SPV_OC_OpAtomicIIncrement,
-      SPV_OC_OpAtomicIDecrement, SPV_OC_OpAtomicIAdd, SPV_OC_OpAtomicISub,
-      SPV_OC_OpAtomicSMin, SPV_OC_OpAtomicUMin, SPV_OC_OpAtomicSMax,
-      SPV_OC_OpAtomicUMax, SPV_OC_OpAtomicAnd, SPV_OC_OpAtomicOr, SPV_OC_OpAtomicXor,
-      SPV_OC_OpPhi, SPV_OC_OpLoopMerge, SPV_OC_OpSelectionMerge, SPV_OC_OpLabel,
-      SPV_OC_OpBranch, SPV_OC_OpBranchConditional, SPV_OC_OpReturn,
-      SPV_OC_OpReturnValue, SPV_OC_OpUnreachable, SPV_OC_OpGroupBroadcast,
-      SPV_OC_OpNoLine, SPV_OC_OpModuleProcessed, SPV_OC_OpGroupNonUniformElect,
+      SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert,
+      SPV_OC_OpPtrCastToGeneric, 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, SPV_OC_OpUnordered, SPV_OC_OpLogicalEqual,
+      SPV_OC_OpLogicalNotEqual, SPV_OC_OpLogicalOr, SPV_OC_OpLogicalAnd,
+      SPV_OC_OpLogicalNot, SPV_OC_OpSelect, SPV_OC_OpIEqual, SPV_OC_OpINotEqual,
+      SPV_OC_OpUGreaterThan, SPV_OC_OpSGreaterThan, SPV_OC_OpUGreaterThanEqual,
+      SPV_OC_OpSGreaterThanEqual, SPV_OC_OpULessThan, SPV_OC_OpSLessThan,
+      SPV_OC_OpULessThanEqual, SPV_OC_OpSLessThanEqual, SPV_OC_OpFOrdEqual,
+      SPV_OC_OpFUnordEqual, SPV_OC_OpFOrdNotEqual, SPV_OC_OpFUnordNotEqual,
+      SPV_OC_OpFOrdLessThan, SPV_OC_OpFUnordLessThan, SPV_OC_OpFOrdGreaterThan,
+      SPV_OC_OpFUnordGreaterThan, SPV_OC_OpFOrdLessThanEqual,
+      SPV_OC_OpFUnordLessThanEqual, SPV_OC_OpFOrdGreaterThanEqual,
+      SPV_OC_OpFUnordGreaterThanEqual, SPV_OC_OpShiftRightLogical,
+      SPV_OC_OpShiftRightArithmetic, SPV_OC_OpShiftLeftLogical, SPV_OC_OpBitwiseOr,
+      SPV_OC_OpBitwiseXor, SPV_OC_OpBitwiseAnd, SPV_OC_OpNot,
+      SPV_OC_OpBitFieldInsert, SPV_OC_OpBitFieldSExtract, SPV_OC_OpBitFieldUExtract,
+      SPV_OC_OpBitReverse, SPV_OC_OpBitCount, SPV_OC_OpControlBarrier,
+      SPV_OC_OpMemoryBarrier, SPV_OC_OpAtomicExchange,
+      SPV_OC_OpAtomicCompareExchange, SPV_OC_OpAtomicCompareExchangeWeak,
+      SPV_OC_OpAtomicIIncrement, SPV_OC_OpAtomicIDecrement, SPV_OC_OpAtomicIAdd,
+      SPV_OC_OpAtomicISub, SPV_OC_OpAtomicSMin, SPV_OC_OpAtomicUMin,
+      SPV_OC_OpAtomicSMax, SPV_OC_OpAtomicUMax, SPV_OC_OpAtomicAnd,
+      SPV_OC_OpAtomicOr, SPV_OC_OpAtomicXor, SPV_OC_OpPhi, SPV_OC_OpLoopMerge,
+      SPV_OC_OpSelectionMerge, SPV_OC_OpLabel, SPV_OC_OpBranch,
+      SPV_OC_OpBranchConditional, SPV_OC_OpReturn, SPV_OC_OpReturnValue,
+      SPV_OC_OpUnreachable, SPV_OC_OpGroupBroadcast, SPV_OC_OpNoLine,
+      SPV_OC_OpModuleProcessed, SPV_OC_OpGroupNonUniformElect,
       SPV_OC_OpGroupNonUniformBroadcast, SPV_OC_OpGroupNonUniformBallot,
       SPV_OC_OpGroupNonUniformIAdd, SPV_OC_OpGroupNonUniformFAdd,
       SPV_OC_OpGroupNonUniformIMul, SPV_OC_OpGroupNonUniformFMul,
       SPV_OC_OpGroupNonUniformSMin, SPV_OC_OpGroupNonUniformUMin,
       SPV_OC_OpGroupNonUniformFMin, SPV_OC_OpGroupNonUniformSMax,
       SPV_OC_OpGroupNonUniformUMax, SPV_OC_OpGroupNonUniformFMax,
       SPV_OC_OpSubgroupBallotKHR, SPV_OC_OpTypeCooperativeMatrixNV,
       SPV_OC_OpCooperativeMatrixLoadNV, SPV_OC_OpCooperativeMatrixStoreNV,

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

class SPV_OCLTernaryArithmeticOp<string mnemonic, int opcode, Type type,
   SPV_OCLTernaryOp<mnemonic, opcode, type, type,
                   traits # [SameOperandsAndResultType]> {
   let assemblyFormat = "operands attr-dict `:` type($result)";
 }
 
 
 // -----
 
-def SPV_OCLFmaOp : SPV_OCLTernaryArithmeticOp<"fma", 26, SPV_Float> {
-  let summary = [{
-    Compute the correctly rounded floating-point representation of the sum
-    of c with the infinitely precise product of a and b. Rounding of
-    intermediate products shall not occur. Edge case results are per the
-    IEEE 754-2008 standard.
-  }];
-
-  let description = [{
-    Result Type, a, b and c must be floating-point or vector(2,3,4,8,16) of
-    floating-point values.
-
-    All of the operands, including the Result Type operand, must be of the
-    same type.
-
-    <!-- End of AutoGen section -->
-
-    ```
-    fma-op ::= ssa-id `=` `spv.OCL.fma` ssa-use, ssa-use, ssa-use `:`
-               float-scalar-vector-type
-    ```mlir
-
-    ```
-    %0 = spv.OCL.fma %a, %b, %c : f32
-    %1 = spv.OCL.fma %a, %b, %c : vector<3xf16>
-    ```
-  }];
-}
-
-// -----
-
 def SPV_OCLCeilOp : SPV_OCLUnaryArithmeticOp<"ceil", 12, SPV_Float> {
   let summary = [{
     Round x to integral value using the round to positive infinity rounding
     mode.
   }];
 
   let description = [{
     Result Type and x must be floating-point or vector(2,3,4,8,16) of
let description = [{
     %2 = spv.OCL.floor %0 : f32
     %3 = spv.OCL.ceifloorl %1 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
+def SPV_OCLFmaOp : SPV_OCLTernaryArithmeticOp<"fma", 26, SPV_Float> {
+  let summary = [{
+    Compute the correctly rounded floating-point representation of the sum
+    of c with the infinitely precise product of a and b. Rounding of
+    intermediate products shall not occur. Edge case results are per the
+    IEEE 754-2008 standard.
+  }];
+
+  let description = [{
+    Result Type, a, b and c must be floating-point or vector(2,3,4,8,16) of
+    floating-point values.
+
+    All of the operands, including the Result Type operand, must be of the
+    same type.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    fma-op ::= ssa-id `=` `spv.OCL.fma` ssa-use, ssa-use, ssa-use `:`
+               float-scalar-vector-type
+    ```mlir
+
+    ```
+    %0 = spv.OCL.fma %a, %b, %c : f32
+    %1 = spv.OCL.fma %a, %b, %c : vector<3xf16>
+    ```
+  }];
+}
+
+// -----
+
 def SPV_OCLLogOp : SPV_OCLUnaryArithmeticOp<"log", 37, SPV_Float> {
   let summary = "Compute the natural logarithm of x.";
 
   let description = [{
     Result Type and x must be floating-point or vector(2,3,4,8,16) of
     floating-point values.
 
     All of the operands, including the Result Type operand, must be of the
let description = [{
     %2 = spv.OCL.tanh %0 : f32
     %3 = spv.OCL.tanh %1 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
+def SPV_PtrCastToGenericOp : SPV_Op<"PtrCastToGeneric", [NoSideEffect]> {
+  let summary = "Convert a pointer’s Storage Class to Generic.";
+
+  let description = [{
+    Result Type must be an OpTypePointer. Its Storage Class must be Generic.
+
+    Pointer must point to the Workgroup, CrossWorkgroup, or Function Storage
+    Class.
+
+    Result Type and Pointer must point to the same type.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    %1 = spv.PtrCastToGeneric %0 : !spv.ptr<vector<3xi64>, CrossWorkgroup> to !spv.ptr<vector<3xi64>, Generic>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[SPV_C_Kernel]>
+  ];
+
+  let arguments = (ins
+    SPV_Type:$pointer
+  );
+
+  let results = (outs
+    SPV_Type:$result
+  );
+}
+
+// -----
+
 def SPV_OCLSAbsOp : SPV_OCLUnaryArithmeticOp<"s_abs", 141, SPV_Integer> {
   let summary = "Absolute value of operand";
 
   let description = [{
     Returns |x|, where x is treated as signed integer.
 
     Result Type and x must be integer or vector(2,3,4,8,16) of
     integer values.

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

 void spirv::PtrAccessChainOp::print(OpAsmPrinter &printer) {
   printAccessChain(*this, concatElemAndIndices(*this), printer);
 }
 
 LogicalResult spirv::PtrAccessChainOp::verify() {
   return verifyAccessChain(*this, indices());
 }
 
+//===----------------------------------------------------------------------===//
+// spv.PtrCastToGenericOp
+//===----------------------------------------------------------------------===//
+
+ParseResult spirv::PtrCastToGenericOp::parse(OpAsmParser &parser,
+                                           OperationState &state) {
+  OpAsmParser::OperandType ptr;
+
+  if (parser.parseOperand(ptr))
+    return failure();
+
+  if (parser.parseColon())
+    return failure();
+
+  Type ptrType;
+  if (parser.parseType(ptrType))
+    return failure();
+
+  if (parser.resolveOperand(ptr, ptrType, state.operands))
+    return failure();
+
+  Type resultType;
+  if (parser.parseKeywordType("to", resultType))
+    return failure();
+
+  state.addTypes(resultType);
+
+  return success();
+}
+
+void spirv::PtrCastToGenericOp::print(OpAsmPrinter &p) {
+  p << " " << getOperand();
+  p << " : " << getOperand().getType() << " to " << getResult().getType();
+}
+
+LogicalResult spirv::PtrCastToGenericOp::verify() {
+  auto opType = getOperand().getType();
+  if (!opType.isa<PointerType>())
+    return failure();
+
+  auto ptrType = opType.cast<PointerType>();
+  if (ptrType.getStorageClass() == StorageClass::Generic)
+    return failure();
+
+  auto resType = getResult().getType();
+  if (!resType.isa<PointerType>())
+    return failure();
+
+  ptrType = resType.cast<PointerType>();
+  if (ptrType.getStorageClass() != StorageClass::Generic)
+    return failure();
+
+  return success();
+}
+
 // TableGen'erated operation interfaces for querying versions, extensions, and
 // capabilities.
 #include "mlir/Dialect/SPIRV/IR/SPIRVAvailability.cpp.inc"
 
 // TablenGen'erated operation definitions.
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.cpp.inc"
 
 namespace mlir {
 namespace spirv {
 // TableGen'erated operation availability interface implementations.
 #include "mlir/Dialect/SPIRV/IR/SPIRVOpAvailabilityImpl.inc"
 } // namespace spirv
 } // namespace mlir

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

 
 // -----
 
 func @fma(%a : vector<3xf32>, %b : vector<3xf32>, %c : vector<3xf32>) -> () {
   // CHECK: spv.OCL.fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : vector<3xf32>
   %2 = spv.OCL.fma %a, %b, %c : vector<3xf32>
   return
 }
+
+//===----------------------------------------------------------------------===//
+// spv.PtrCastToGenericOp
+//===----------------------------------------------------------------------===//
+
+func @cast_to_generic(%a: !spv.ptr<i32, CrossWorkgroup>) {
+  // CHECKL spv.PtrCastToGeneric {{%[^,]*}} : !spv.ptr<i32, CrossWorkgroup> to %!spv.ptr<i32, Generic>
+  %0 = spv.PtrCastToGeneric %a : !spv.ptr<i32, CrossWorkgroup> to %!spv.ptr<i32, Generic>
+}