[mlir][spirv] Add casting ops to/from generic storage space

[antiagainst]

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

336

These ops should not derive from SPV_CastOp; they should derive from SPV_Op directly. (SPV_CastOp is meant for numeric value casting; here we are doing pointer casting.)

360

let arguments = (ins
  SPV_AnyPtr:$pointer
);

let results = (outs
  SPV_AnyPtr:$result
);

let assemblyFormat = [{
  $operand attr-dict `:` type($ptr) `to` type($result)
}];

362

This is missing availability spec (which says these ops depend on Kernel)?

Did you pull these op definitions in using https://github.com/llvm/llvm-project/blob/main/mlir/utils/spirv/define_inst.sh?

366

Similarly like the above op.

396

Similarly like the above op.

423

Note that for this op the spec requires two operands, one pointer, one storage. We have two choices:

1. Define the op as what the spec requires---two operands. The storage operands kinda duplicates with the storage class in the result type. Pro is that we don't need to manually code the serialization/deserialization.
2. Define the op as only having one operand. This would make the IR cleaner. It would require customized (de)serialization though.

My preference is 2. (You can find examples of how to write customized (de)serialization in the SerializeOps.cpp or DeseralizeOps.cpp.) But if you want to go with 1 i'm also fine with it.

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

1535

This check isn't necessary if we have let arguments = (ins SPV_AnyPtr:$pointer) as said in the above.

1569

Same here.

1603

Same here.

1618

We need to additionally verify the storage operand matches the storage class in the result type (if defining the op as taking two operands).

[nirvedhmeshram]

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

423

Took a stab at option 2.

[antiagainst]

Nice, thanks! I just have a few more nits. Feel free to land after addressing them. :)

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

364

Please name this as $pointer to follow the spec.

Similarly for the other two ops.

mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp

552 ↗

(On Diff #461698)

Use opBuilder.create<spirv::GenericCastToPtrExplicitOp>(..) to avoid the explicit OperationState and string name?

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_OpGenericCastToPtr           : I32EnumAttrCase<"OpGenericCastToPtr", 122>;
+def SPV_OC_OpGenericCastToPtrExplicit   : I32EnumAttrCase<"OpGenericCastToPtrExplicit", 123>;
 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", "opcode", [
       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_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert, SPV_OC_OpPtrCastToGeneric, 
+      SPV_OC_OpGenericCastToPtr, SPV_OC_OpGenericCastToPtrExplicit, 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_OpVectorTimesScalar,
       SPV_OC_OpMatrixTimesScalar, SPV_OC_OpMatrixTimesMatrix, SPV_OC_OpIAddCarry,
       SPV_OC_OpISubBorrow, 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,

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

let description = [{
 
     ```mlir
     %1 = spv.UConvertOp %0 : i32 to i64
     %3 = spv.UConvertOp %2 : vector<3xi32> to vector<3xi64>
     ```
   }];
 }
 
+// -----
+
+def SPV_PtrCastToGenericOp : SPV_CastOp<"PtrCastToGeneric",
+                                SPV_AnyPtr,
+                                SPV_AnyPtr,
+                                [NoSideEffect]> {
+  let summary = [{
+    Convert pointer to Generic Storage.
+  }];
+
+  let description = [{
+    Convert a pointer’s Storage Class to Generic.
+
+    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.
+    
+    #### Example:
+
+    ```mlir
+    %1 = spv.PtrCastToGenericOp %0 : !spv.ptr<f32, CrossWorkGroup> to 
+         !spv.ptr<f32, Generic>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GenericCastToPtrOp : SPV_CastOp<"GenericCastToPtr",
+                                SPV_AnyPtr,
+                                SPV_AnyPtr,
+                                [NoSideEffect]> {
+  let summary = [{
+    Convert pointer from  Generic Storage to Workgroup, CrossWorkgroup,
+    or Function.
+  }];
+
+  let description = [{
+    Convert a pointer’s Storage Class to a non-Generic class.
+
+    Result Type must be an OpTypePointer. Its Storage Class 
+    must be Workgroup, CrossWorkgroup, or Function.
+
+    Pointer must point to the Generic Storage Class.
+
+    Result Type and Pointer must point to the same type.
+
+    #### Example:
+
+    ```mlir
+       %1 = spv.GenericCastToPtrOp %0 : !spv.ptr<f32, Generic> to
+       !spv.ptr<f32, CrossWorkGroup>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GenericCastToPtrExplicitOp : SPV_CastOp<"GenericCastToPtrExplicit",
+                                SPV_AnyPtr,
+                                SPV_AnyPtr,
+                                [NoSideEffect]> {
+  let summary = [{
+    Convert pointer from  Generic Storage to Workgroup, CrossWorkgroup,
+    or Function explicitly.
+  }];
+
+  let description = [{
+    Attempts to explicitly convert Pointer to non-Generic 
+    storage-class pointer value.
+
+    Result Type must be an OpTypePointer. Its Storage Class must be
+    Workgroup, CrossWorkgroup, or Function.
+
+    Pointer must have a type of OpTypePointer whose Type is the same as 
+    the Type of Result Type. Pointer must point to the Generic Storage Class.
+    If the cast fails, the instruction result is an OpConstantNull pointer 
+    in the Storage Class of result Type.
+
+    #### Example:
+
+    ```mlir
+       %1 = spv.GenericCastToPtrExplicitOp %0 : !spv.ptr<f32, Generic> to
+       !spv.ptr<f32, CrossWorkGroup>
+    ```
+  }];
+}
 #endif // MLIR_DIALECT_SPIRV_IR_CAST_OPS

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

if (operandBitWidth != resultBitWidth) {
     return emitOpError("mismatch in result type bitwidth ")
            << resultBitWidth << " and operand type bitwidth "
            << operandBitWidth;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// spv.PtrCastToGenericOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult spirv::PtrCastToGenericOp::verify() {
+  auto operandType = operand().getType();
+  auto resultType = result().getType();
+  if (!operandType.isa<spirv::PointerType>() ||
+      !resultType.isa<spirv::PointerType>()) {
+    return emitError("operand and result expected to be pointer type");
+  }
+  spirv::StorageClass operandStorage =
+      operandType.cast<spirv::PointerType>().getStorageClass();
+  if (operandStorage != spirv::StorageClass::Workgroup &&
+      operandStorage != spirv::StorageClass::CrossWorkgroup &&
+      operandStorage != spirv::StorageClass::Function)
+    return emitError("pointer must point to the Workgroup, CrossWorkgroup"
+                     ", or Function Storage Class");
+  spirv::StorageClass resultStorage =
+      resultType.cast<spirv::PointerType>().getStorageClass();
+  if (resultStorage != spirv::StorageClass::Generic)
+    return emitError("result type must be of storage class Generic");
+
+  Type operandPointeeType =
+      operandType.cast<spirv::PointerType>().getPointeeType();
+  Type resultPointeeType =
+      resultType.cast<spirv::PointerType>().getPointeeType();
+  if (operandPointeeType != resultPointeeType)
+    return emitOpError("pointer operand's pointee type must have the same "
+                       "as the op result type, but found ")
+           << operandPointeeType << " vs " << resultPointeeType;
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GenericCastToPtrOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult spirv::GenericCastToPtrOp::verify() {
+  auto operandType = operand().getType();
+  auto resultType = result().getType();
+  if (!operandType.isa<spirv::PointerType>() ||
+      !resultType.isa<spirv::PointerType>()) {
+    return emitError("operand and result expected to be pointer type");
+  }
+  spirv::StorageClass operandStorage =
+      operandType.cast<spirv::PointerType>().getStorageClass();
+  if (operandStorage != spirv::StorageClass::Generic)
+    return emitError("pointer type must be of storage class Generic");
+  spirv::StorageClass resultStorage =
+      resultType.cast<spirv::PointerType>().getStorageClass();
+  if (resultStorage != spirv::StorageClass::Workgroup &&
+      resultStorage != spirv::StorageClass::CrossWorkgroup &&
+      resultStorage != spirv::StorageClass::Function)
+    return emitError("result must point to the Workgroup, CrossWorkgroup, "
+                     "or Function Storage Class");
+
+  Type operandPointeeType =
+      operandType.cast<spirv::PointerType>().getPointeeType();
+  Type resultPointeeType =
+      resultType.cast<spirv::PointerType>().getPointeeType();
+  if (operandPointeeType != resultPointeeType)
+    return emitOpError("pointer operand's pointee type must have the same "
+                       "as the op result type, but found ")
+           << operandPointeeType << " vs " << resultPointeeType;
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GenericCastToPtrExplicitOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult spirv::GenericCastToPtrExplicitOp::verify() {
+  auto operandType = operand().getType();
+  auto resultType = result().getType();
+  if (!operandType.isa<spirv::PointerType>() ||
+      !resultType.isa<spirv::PointerType>()) {
+    return emitError("operand and result expected to be pointer type");
+  }
+  spirv::StorageClass operandStorage =
+      operandType.cast<spirv::PointerType>().getStorageClass();
+  if (operandStorage != spirv::StorageClass::Generic)
+    return emitError("pointer type must be of storage class Generic");
+  spirv::StorageClass resultStorage =
+      resultType.cast<spirv::PointerType>().getStorageClass();
+  if (resultStorage != spirv::StorageClass::Workgroup &&
+      resultStorage != spirv::StorageClass::CrossWorkgroup &&
+      resultStorage != spirv::StorageClass::Function)
+    return emitError("result must point to the Workgroup, CrossWorkgroup, "
+                     "or Function Storage Class");
+
+  Type operandPointeeType =
+      operandType.cast<spirv::PointerType>().getPointeeType();
+  Type resultPointeeType =
+      resultType.cast<spirv::PointerType>().getPointeeType();
+  if (operandPointeeType != resultPointeeType)
+    return emitOpError("pointer operand's pointee type must have the same "
+                       "as the op result type, but found ")
+           << operandPointeeType << " vs " << resultPointeeType;
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // spv.BranchOp
 //===----------------------------------------------------------------------===//
 
 SuccessorOperands spirv::BranchOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return SuccessorOperands(0, targetOperandsMutable());
 }
 

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

 //===----------------------------------------------------------------------===//
 
 func.func @u_convert_scalar(%arg0 : i32) -> i64 {
   // CHECK: {{%.*}} = spv.UConvert {{%.*}} : i32 to i64
   %0 = spv.UConvert %arg0 : i32 to i64
   spv.ReturnValue %0 : i64
 }
 
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.PtrCastToGeneric
+//===----------------------------------------------------------------------===//
+
+func.func @ptrcasttogeneric1(%arg0 : !spv.ptr<f32, CrossWorkgroup>) {
+  // CHECK: {{%.*}} = spv.PtrCastToGeneric {{%.*}} : !spv.ptr<f32, CrossWorkgroup> to !spv.ptr<f32, Generic>
+  %0 = spv.PtrCastToGeneric %arg0 : !spv.ptr<f32, CrossWorkgroup> to !spv.ptr<f32, Generic>
+  return
+}
+// -----
+
+func.func @ptrcasttogeneric2(%arg0 : !spv.ptr<f32, StorageBuffer>) {
+  // expected-error @+1 {{pointer must point to the Workgroup, CrossWorkgroup, or Function Storage Class}}
+  %0 = spv.PtrCastToGeneric %arg0 : !spv.ptr<f32, StorageBuffer> to !spv.ptr<f32, Generic>
+  return
+}
+
+// -----
+
+func.func @ptrcasttogeneric3(%arg0 : !spv.ptr<f32, CrossWorkgroup>) {
+  // expected-error @+1 {{result type must be of storage class Generic}}
+  %0 = spv.PtrCastToGeneric %arg0 : !spv.ptr<f32, CrossWorkgroup> to !spv.ptr<f32, Uniform>
+  return
+}
+
+// -----
+
+func.func @ptrcasttogeneric4(%arg0 : !spv.ptr<f32, CrossWorkgroup>) {
+  // expected-error @+1 {{pointee type must have the same as the op result type}}
+  %0 = spv.PtrCastToGeneric %arg0 : !spv.ptr<f32, CrossWorkgroup> to !spv.ptr<vector<2xi32>, Generic>
+  return
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GenericCastToPtr
+//===----------------------------------------------------------------------===//
+
+func.func @genericcasttoptr1(%arg0 : !spv.ptr<vector<2xi32>, Generic>) {
+  // CHECK: {{%.*}} = spv.GenericCastToPtr {{%.*}} : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+  %0 = spv.GenericCastToPtr %arg0 : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+  return
+}
+// -----
+
+func.func @genericcasttoptr2(%arg0 : !spv.ptr<f32, Uniform>) {
+  // expected-error @+1 {{pointer type must be of storage class Generic}}
+  %0 = spv.GenericCastToPtr %arg0 : !spv.ptr<f32, Uniform> to !spv.ptr<f32, Workgroup>
+  return
+}
+
+// -----
+
+func.func @genericcasttoptr3(%arg0 : !spv.ptr<f32, Generic>) {
+  // expected-error @+1 {{result must point to the Workgroup, CrossWorkgroup, or Function Storage Class}}
+  %0 = spv.GenericCastToPtr %arg0 : !spv.ptr<f32, Generic> to !spv.ptr<f32, Uniform>
+  return
+}
+
+// -----
+
+func.func @genericcasttoptr4(%arg0 : !spv.ptr<f32, Generic>) {
+  // expected-error @+1 {{pointee type must have the same as the op result type}}
+  %0 = spv.GenericCastToPtr %arg0 : !spv.ptr<f32, Generic> to !spv.ptr<vector<2xi32>, Workgroup>
+  return
+}
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GenericCastToPtrExplicit
+//===----------------------------------------------------------------------===//
+
+func.func @genericcasttoptrexplicit1(%arg0 : !spv.ptr<vector<2xi32>, Generic>) {
+  // CHECK: {{%.*}} = spv.GenericCastToPtrExplicit {{%.*}} : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+  %0 = spv.GenericCastToPtrExplicit %arg0 : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+  return
+}
+// -----
+
+func.func @genericcasttoptrexplicit2(%arg0 : !spv.ptr<f32, Uniform>) {
+  // expected-error @+1 {{pointer type must be of storage class Generic}}
+  %0 = spv.GenericCastToPtrExplicit %arg0 : !spv.ptr<f32, Uniform> to !spv.ptr<f32, Workgroup>
+  return
+}
+
+// -----
+
+func.func @genericcasttoptrexplicit3(%arg0 : !spv.ptr<f32, Generic>) {
+  // expected-error @+1 {{result must point to the Workgroup, CrossWorkgroup, or Function Storage Class}}
+  %0 = spv.GenericCastToPtrExplicit %arg0 : !spv.ptr<f32, Generic> to !spv.ptr<f32, Uniform>
+  return
+}
+
+// -----
+
+func.func @genericcasttoptrexplicit4(%arg0 : !spv.ptr<f32, Generic>) {
+  // expected-error @+1 {{pointee type must have the same as the op result type}}
+  %0 = spv.GenericCastToPtrExplicit %arg0 : !spv.ptr<f32, Generic> to !spv.ptr<vector<2xi32>, Workgroup>
+  return
+}

mlir/test/Target/SPIRV/cast-ops.mlir

spv.func @s_convert(%arg0 : i32) -> i64 "None" {
     spv.ReturnValue %0 : i64
   }
   spv.func @u_convert(%arg0 : i32) -> i64 "None" {
     // CHECK: {{%.*}} = spv.UConvert {{%.*}} : i32 to i64
     %0 = spv.UConvert %arg0 : i32 to i64
     spv.ReturnValue %0 : i64
   }
 }
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Kernel], []> {
+  spv.func @ptr_cast_to_generic(%arg0 : !spv.ptr<f32, CrossWorkgroup>) "None" {
+    // CHECK: {{%.*}} = spv.PtrCastToGeneric {{%.*}} : !spv.ptr<f32, CrossWorkgroup> to !spv.ptr<f32, Generic>
+    %0 = spv.PtrCastToGeneric %arg0 : !spv.ptr<f32, CrossWorkgroup> to !spv.ptr<f32, Generic>
+    spv.Return
+  }
+  spv.func @generic_cast_to_ptr(%arg0 : !spv.ptr<vector<2xi32>, Generic>) "None" {
+    // CHECK: {{%.*}} = spv.GenericCastToPtr {{%.*}} : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+    %0 = spv.GenericCastToPtr %arg0 : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+    spv.Return
+  }
+  spv.func @generic_cast_to_ptr_explicit(%arg0 : !spv.ptr<vector<2xi32>, Generic>) "None" {
+    // CHECK: {{%.*}} = spv.GenericCastToPtrExplicit {{%.*}} : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+    %0 = spv.GenericCastToPtrExplicit %arg0 : !spv.ptr<vector<2xi32>, Generic> to !spv.ptr<vector<2xi32>, CrossWorkgroup>
+    spv.Return
+  }
+}