[mlir][spirv] Add block read and write from SPV_INTEL_subgroups

[abialas]

Summary

Added support to OpSubgroupBlockReadINTEL and OpSubgroupBlockWriteINTEL

[antiagainst]

Nice, thanks!

[ThomasRaoux]

mlir/include/mlir/Dialect/SPIRV/SPIRVGroupOps.td

228

Are you missing the buffer type? (ex: "StorageBuffer")

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2056

The description mentions:

Result Type may be a scalar or vector type, and its component type must be
equal to the type pointed to by Ptr.

The type of Ptr must be a pointer type, and must point to a scalar type.

Here it looks like if the data is a vector we would expect the pointer to be a pointer of elementType which would be a pointer of vector.

2074–2077

Same problem as above here, this verification function would check that pointer points to a type matching the destination type. This is consistent with the parser but not with the description.

It might be worth adding a test for the vector case to make sure you have what you want.

[abialas]

added changes from code review

[abialas]

Thanks Thomas, good catch. Please take a look at the updated version. If it looks good, may I ask you to push this change to llvm master?

[ThomasRaoux]

In D86876#2251655, @abialas wrote:

Thanks Thomas, good catch. Please take a look at the updated version. If it looks good, may I ask you to push this change to llvm master?

Sure, this is done.

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

 def SPV_OC_OpGroupNonUniformUMax       : I32EnumAttrCase<"OpGroupNonUniformUMax", 357>;
 def SPV_OC_OpGroupNonUniformFMax       : I32EnumAttrCase<"OpGroupNonUniformFMax", 358>;
 def SPV_OC_OpSubgroupBallotKHR         : I32EnumAttrCase<"OpSubgroupBallotKHR", 4421>;
 def SPV_OC_OpTypeCooperativeMatrixNV   : I32EnumAttrCase<"OpTypeCooperativeMatrixNV", 5358>;
 def SPV_OC_OpCooperativeMatrixLoadNV   : I32EnumAttrCase<"OpCooperativeMatrixLoadNV", 5359>;
 def SPV_OC_OpCooperativeMatrixStoreNV  : I32EnumAttrCase<"OpCooperativeMatrixStoreNV", 5360>;
 def SPV_OC_OpCooperativeMatrixMulAddNV : I32EnumAttrCase<"OpCooperativeMatrixMulAddNV", 5361>;
 def SPV_OC_OpCooperativeMatrixLengthNV : I32EnumAttrCase<"OpCooperativeMatrixLengthNV", 5362>;
+def SPV_OC_OpSubgroupBlockReadINTEL    : I32EnumAttrCase<"OpSubgroupBlockReadINTEL", 5575>;
+def SPV_OC_OpSubgroupBlockWriteINTEL   : I32EnumAttrCase<"OpSubgroupBlockWriteINTEL", 5576>;
 
 def SPV_OpcodeAttr :
     SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpNop, SPV_OC_OpUndef, SPV_OC_OpSourceContinued, SPV_OC_OpSource,
       SPV_OC_OpSourceExtension, SPV_OC_OpName, SPV_OC_OpMemberName, SPV_OC_OpString,
       SPV_OC_OpLine, SPV_OC_OpExtension, SPV_OC_OpExtInstImport, SPV_OC_OpExtInst,
       SPV_OC_OpMemoryModel, SPV_OC_OpEntryPoint, SPV_OC_OpExecutionMode,
       SPV_OC_OpCapability, SPV_OC_OpTypeVoid, SPV_OC_OpTypeBool, SPV_OC_OpTypeInt,
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       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, SPV_OC_OpCooperativeMatrixMulAddNV,
-      SPV_OC_OpCooperativeMatrixLengthNV
+      SPV_OC_OpCooperativeMatrixLengthNV, SPV_OC_OpSubgroupBlockReadINTEL,
+      SPV_OC_OpSubgroupBlockWriteINTEL
     ]>;
 
 // End opcode section. Generated from SPIR-V spec; DO NOT MODIFY!
 
 //===----------------------------------------------------------------------===//
 // SPIR-V op definitions
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Dialect/SPIRV/SPIRVGroupOps.td

def SPV_GroupBroadcastOp : SPV_Op<"GroupBroadcast",
 
   let results = (outs
     SPV_Type:$result
   );
 
   let assemblyFormat = [{
     $execution_scope operands attr-dict `:` type($value) `,` type($localid)
   }];
-
 }
 
 // -----
 
 def SPV_SubgroupBallotKHROp : SPV_Op<"SubgroupBallotKHR", []> {
   let summary = "See extension SPV_KHR_shader_ballot";
 
   let description = [{
def SPV_SubgroupBallotKHROp : SPV_Op<"SubgroupBallotKHR", []> {
 
   let verifier = [{ return success(); }];
 
   let assemblyFormat = "$predicate attr-dict `:` type($result)";
 }
 
 // -----
 
+def SPV_SubgroupBlockReadINTELOp : SPV_Op<"SubgroupBlockReadINTEL", []> {
+  let summary = "See extension SPV_INTEL_subgroups";
+
+  let description = [{
+    Reads one or more components of Result data for each invocation in the
+    subgroup from the specified Ptr as a block operation.
+
+    The data is read strided, so the first value read is:
+    Ptr[ SubgroupLocalInvocationId ]
+
+    and the second value read is:
+    Ptr[ SubgroupLocalInvocationId + SubgroupMaxSize ]
+    etc.
+
+    Result Type may be a scalar or vector type, and its component type must be
+    equal to the type pointed to by Ptr.
+
+    The type of Ptr must be a pointer type, and must point to a scalar type.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    subgroup-block-read-INTEL-op ::= ssa-id `=` `spv.SubgroupBlockReadINTEL`
+                                storage-class ssa_use `:` spirv-element-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    %0 = spv.SubgroupBlockReadINTEL "StorageBuffer" %ptr : i32
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[SPV_INTEL_subgroups]>,
+    Capability<[SPV_C_SubgroupBufferBlockIOINTEL]>
+  ];
+
+  let arguments = (ins
+    SPV_AnyPtr:$ptr
+  );
+
+  let results = (outs
+    SPV_Type:$value
+  );
+}
+
+// -----
+
+def SPV_SubgroupBlockWriteINTELOp : SPV_Op<"SubgroupBlockWriteINTEL", []> {
+  let summary = "See extension SPV_INTEL_subgroups";
+
+  let description = [{
+    Writes one or more components of Data for each invocation in the subgroup
+    from the specified Ptr as a block operation.
+
+    The data is written strided, so the first value is written to:
+    Ptr[ SubgroupLocalInvocationId ]
+
+    and the second value written is:
+    Ptr[ SubgroupLocalInvocationId + SubgroupMaxSize ]
+    etc.
+
+    The type of Ptr must be a pointer type, and must point to a scalar type.
+
+    The component type of Data must be equal to the type pointed to by Ptr.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    subgroup-block-write-INTEL-op ::= ssa-id `=` `spv.SubgroupBlockWriteINTEL`
+                      storage-class ssa_use `,` ssa-use `:` spirv-element-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    spv.SubgroupBlockWriteINTEL "StorageBuffer" %ptr, %value : i32
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[SPV_INTEL_subgroups]>,
+    Capability<[SPV_C_SubgroupBufferBlockIOINTEL]>
+  ];
+
+  let arguments = (ins
+    SPV_AnyPtr:$ptr,
+    SPV_Type:$value
+  );
+
+  let results = (outs);
+}
+
+// -----
+
 #endif // SPIRV_GROUP_OPS

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

static LogicalResult verifyLoadStorePtrAndValTypes(LoadStoreOpTy op, Value ptr,
   // SPIR-V OpLoad/OpStore operations
   if (val.getType() !=
       ptr.getType().cast<spirv::PointerType>().getPointeeType()) {
     return op.emitOpError("mismatch in result type and pointer type");
   }
   return success();
 }
 
+template <typename BlockReadWriteOpTy>
+static LogicalResult verifyBlockReadWritePtrAndValTypes(BlockReadWriteOpTy op,
+                                                        Value ptr, Value val) {
+  auto valType = val.getType();
+  if (auto valVecTy = valType.dyn_cast<VectorType>())
+    valType = valVecTy.getElementType();
+
+  if (valType != ptr.getType().cast<spirv::PointerType>().getPointeeType()) {
+    return op.emitOpError("mismatch in result type and pointer type");
+  }
+  return success();
+}
+
 static ParseResult parseVariableDecorations(OpAsmParser &parser,
                                             OperationState &state) {
   auto builtInName = llvm::convertToSnakeFromCamelCase(
       stringifyDecoration(spirv::Decoration::BuiltIn));
   if (succeeded(parser.parseOptionalKeyword("bind"))) {
     Attribute set, binding;
     // Parse optional descriptor binding
     auto descriptorSetName = llvm::convertToSnakeFromCamelCase(
static LogicalResult verify(spirv::GroupNonUniformBallotOp ballotOp) {
   if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
     return ballotOp.emitOpError(
         "execution scope must be 'Workgroup' or 'Subgroup'");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// spv.SubgroupBlockReadINTEL
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseSubgroupBlockReadINTELOp(OpAsmParser &parser,
+                                                 OperationState &state) {
+  // Parse the storage class specification
+  spirv::StorageClass storageClass;
+  OpAsmParser::OperandType ptrInfo;
+  Type elementType;
+  if (parseEnumStrAttr(storageClass, parser) || parser.parseOperand(ptrInfo) ||
+      parser.parseColon() || parser.parseType(elementType)) {
+    return failure();
+  }
+
+  auto ptrType = spirv::PointerType::get(elementType, storageClass);
+  if (auto valVecTy = elementType.dyn_cast<VectorType>())
+    ptrType = spirv::PointerType::get(valVecTy.getElementType(), storageClass);
+
+  if (parser.resolveOperand(ptrInfo, ptrType, state.operands)) {
+    return failure();
+  }
+
+  state.addTypes(elementType);
+  return success();
+}
+
+static void print(spirv::SubgroupBlockReadINTELOp blockReadOp,
+                  OpAsmPrinter &printer) {
+  SmallVector<StringRef, 4> elidedAttrs;
+  printer << spirv::SubgroupBlockReadINTELOp::getOperationName() << " "
+          << blockReadOp.ptr();
+  printer << " : " << blockReadOp.getType();
+}
+
+static LogicalResult verify(spirv::SubgroupBlockReadINTELOp blockReadOp) {
+  if (failed(verifyBlockReadWritePtrAndValTypes(blockReadOp, blockReadOp.ptr(),
+                                                blockReadOp.value())))
+    return failure();
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// spv.SubgroupBlockWriteINTEL
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseSubgroupBlockWriteINTELOp(OpAsmParser &parser,
+                                                  OperationState &state) {
+  // Parse the storage class specification
+  spirv::StorageClass storageClass;
+  SmallVector<OpAsmParser::OperandType, 2> operandInfo;
+  auto loc = parser.getCurrentLocation();
+  Type elementType;
+  if (parseEnumStrAttr(storageClass, parser) ||
+      parser.parseOperandList(operandInfo, 2) || parser.parseColon() ||
+      parser.parseType(elementType)) {
+    return failure();
+  }
+
+  auto ptrType = spirv::PointerType::get(elementType, storageClass);
+  if (auto valVecTy = elementType.dyn_cast<VectorType>())
+    ptrType = spirv::PointerType::get(valVecTy.getElementType(), storageClass);
+
+  if (parser.resolveOperands(operandInfo, {ptrType, elementType}, loc,
+                             state.operands)) {
+    return failure();
+  }
+  return success();
+}
+
+static void print(spirv::SubgroupBlockWriteINTELOp blockWriteOp,
+                  OpAsmPrinter &printer) {
+  SmallVector<StringRef, 4> elidedAttrs;
+  printer << spirv::SubgroupBlockWriteINTELOp::getOperationName() << " "
+          << blockWriteOp.ptr() << ", " << blockWriteOp.value();
+  printer << " : " << blockWriteOp.value().getType();
+}
+
+static LogicalResult verify(spirv::SubgroupBlockWriteINTELOp blockWriteOp) {
+  if (failed(verifyBlockReadWritePtrAndValTypes(
+          blockWriteOp, blockWriteOp.ptr(), blockWriteOp.value())))
+    return failure();
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // spv.GroupNonUniformElectOp
 //===----------------------------------------------------------------------===//
 
 void spirv::GroupNonUniformElectOp::build(OpBuilder &builder,
                                           OperationState &state,
                                           spirv::Scope scope) {
   build(builder, state, builder.getI1Type(), scope);
 }

mlir/test/Dialect/SPIRV/Serialization/group-ops.mlir

spv.func @group_broadcast_1(%value: f32, %localid: i32 ) -> f32 "None" {
     spv.ReturnValue %0: f32
   }
   // CHECK-LABEL: @group_broadcast_2
   spv.func @group_broadcast_2(%value: f32, %localid: vector<3xi32> ) -> f32 "None" {
     // CHECK: spv.GroupBroadcast "Workgroup" %{{.*}}, %{{.*}} : f32, vector<3xi32>
     %0 = spv.GroupBroadcast "Workgroup" %value, %localid : f32, vector<3xi32>
     spv.ReturnValue %0: f32
   }
+  // CHECK-LABEL: @subgroup_block_read_intel
+  spv.func @subgroup_block_read_intel(%ptr : !spv.ptr<i32, StorageBuffer>) -> i32 "None" {
+    // CHECK: spv.SubgroupBlockReadINTEL %{{.*}} : i32
+    %0 = spv.SubgroupBlockReadINTEL "StorageBuffer" %ptr : i32
+    spv.ReturnValue %0: i32
+  }
+  // CHECK-LABEL: @subgroup_block_read_intel_vector
+  spv.func @subgroup_block_read_intel_vector(%ptr : !spv.ptr<i32, StorageBuffer>) -> vector<3xi32> "None" {
+    // CHECK: spv.SubgroupBlockReadINTEL %{{.*}} : vector<3xi32>
+    %0 = spv.SubgroupBlockReadINTEL "StorageBuffer" %ptr : vector<3xi32>
+    spv.ReturnValue %0: vector<3xi32>
+  }
+  // CHECK-LABEL: @subgroup_block_write_intel
+  spv.func @subgroup_block_write_intel(%ptr : !spv.ptr<i32, StorageBuffer>, %value: i32) -> () "None" {
+    // CHECK: spv.SubgroupBlockWriteINTEL %{{.*}}, %{{.*}} : i32
+    spv.SubgroupBlockWriteINTEL "StorageBuffer" %ptr, %value : i32
+    spv.Return
+  }
+  // CHECK-LABEL: @subgroup_block_write_intel_vector
+  spv.func @subgroup_block_write_intel_vector(%ptr : !spv.ptr<i32, StorageBuffer>, %value: vector<3xi32>) -> () "None" {
+    // CHECK: spv.SubgroupBlockWriteINTEL %{{.*}}, %{{.*}} : vector<3xi32>
+    spv.SubgroupBlockWriteINTEL "StorageBuffer" %ptr, %value : vector<3xi32>
+    spv.Return
+  }
 }

mlir/test/Dialect/SPIRV/group-ops.mlir

 
 // -----
 
 func @group_broadcast_negative_locid_vec4(%value: f32, %localid: vector<4xi32> ) -> f32 {
   // expected-error @+1 {{localid is a vector and can be with only  2 or 3 components, actual number is 4}}
   %0 = spv.GroupBroadcast "Subgroup" %value, %localid : f32, vector<4xi32>
   return %0: f32
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.SubgroupBlockReadINTEL
+//===----------------------------------------------------------------------===//
+
+func @subgroup_block_read_intel(%ptr : !spv.ptr<i32, StorageBuffer>) -> i32 {
+  // CHECK: spv.SubgroupBlockReadINTEL %{{.*}} : i32
+  %0 = spv.SubgroupBlockReadINTEL "StorageBuffer" %ptr : i32
+  return %0: i32
+}
+
+// -----
+
+func @subgroup_block_read_intel_vector(%ptr : !spv.ptr<i32, StorageBuffer>) -> vector<3xi32> {
+  // CHECK: spv.SubgroupBlockReadINTEL %{{.*}} : vector<3xi32>
+  %0 = spv.SubgroupBlockReadINTEL "StorageBuffer" %ptr : vector<3xi32>
+  return %0: vector<3xi32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.SubgroupBlockWriteINTEL
+//===----------------------------------------------------------------------===//
+
+func @subgroup_block_write_intel(%ptr : !spv.ptr<i32, StorageBuffer>, %value: i32) -> () {
+  // CHECK: spv.SubgroupBlockWriteINTEL %{{.*}}, %{{.*}} : i32
+  spv.SubgroupBlockWriteINTEL "StorageBuffer" %ptr, %value : i32
+  return
+}
+
+// -----
+
+func @subgroup_block_write_intel_vector(%ptr : !spv.ptr<i32, StorageBuffer>, %value: vector<3xi32>) -> () {
+  // CHECK: spv.SubgroupBlockWriteINTEL %{{.*}}, %{{.*}} : vector<3xi32>
+  spv.SubgroupBlockWriteINTEL "StorageBuffer" %ptr, %value : vector<3xi32>
+  return
+}
 No newline at end of file