[mlir][spirv] add AtomicFAddEXTOp

[antiagainst]

Thanks! Just a few nits.

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

288

Could you provide some documentation for this op after the AutoGen section?

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

760

Nit: StringifyTypeName?

[Hardcode84]

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

288

Copypasted from https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/EXT/SPV_EXT_shader_atomic_float_add.asciidoc

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

 
 #ifndef MLIR_DIALECT_SPIRV_IR_ATOMIC_OPS
 #define MLIR_DIALECT_SPIRV_IR_ATOMIC_OPS
 
 class SPV_AtomicUpdateOp<string mnemonic, list<OpTrait> traits = []> :
   SPV_Op<mnemonic, traits> {
   let parser = [{ return ::parseAtomicUpdateOp(parser, result, false); }];
   let printer = [{ return ::printAtomicUpdateOp(getOperation(), p); }];
-  let verifier = [{ return ::verifyAtomicUpdateOp(getOperation()); }];
+  let verifier = [{ return ::verifyAtomicUpdateOp<IntegerType>(getOperation()); }];
 
   let arguments = (ins
     SPV_AnyPtr:$pointer,
     SPV_ScopeAttr:$memory_scope,
     SPV_MemorySemanticsAttr:$semantics
   );
 
   let results = (outs
     SPV_Integer:$result
   );
 }
 
 class SPV_AtomicUpdateWithValueOp<string mnemonic, list<OpTrait> traits = []> :
   SPV_Op<mnemonic, traits> {
   let parser = [{ return ::parseAtomicUpdateOp(parser, result, true); }];
   let printer = [{ return ::printAtomicUpdateOp(getOperation(), p); }];
-  let verifier = [{ return ::verifyAtomicUpdateOp(getOperation()); }];
+  let verifier = [{ return ::verifyAtomicUpdateOp<IntegerType>(getOperation()); }];
 
   let arguments = (ins
     SPV_AnyPtr:$pointer,
     SPV_ScopeAttr:$memory_scope,
     SPV_MemorySemanticsAttr:$semantics,
     SPV_Integer:$value
   );
 
def SPV_AtomicExchangeOp : SPV_Op<"AtomicExchange", []> {
 
   let results = (outs
     SPV_Numerical:$result
   );
 }
 
 // -----
 
+def SPV_AtomicFAddEXTOp : SPV_Op<"AtomicFAddEXT", []> {
+  let summary = "TBD";
+
+  let description = [{
+
+
+    <!-- End of AutoGen section -->
+
+    Perform the following steps atomically with respect to any other atomic
+    accesses within Scope to the same location:
+
+    1) load through Pointer to get an Original Value,
+
+    2) get a New Value by float addition of Original Value and Value, and
+
+    3) store the New Value back through Pointer.
+
+    The instruction’s result is the Original Value.
+
+    Result Type must be a floating-point type scalar.
+
+    The type of Value must be the same as Result Type. The type of the value
+    pointed to by Pointer must be the same as Result Type.
+
+    Memory must be a valid memory Scope.
+
+    ```
+    atomic-fadd-op ::=
+        `spv.AtomicFAddEXT` scope memory-semantics
+                            ssa-use `,` ssa-use `:` spv-pointer-type
+    ```
+
+    #### Example:
+
+    ```mlir
+    %0 = spv.AtomicFAddEXT "Device" "None" %pointer, %value :
+                           !spv.ptr<f32, StorageBuffer>
+    ```mlir
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[SPV_EXT_shader_atomic_float_add]>,
+    Capability<[SPV_C_AtomicFloat16AddEXT, SPV_C_AtomicFloat32AddEXT, SPV_C_AtomicFloat64AddEXT]>
+  ];
+
+  let arguments = (ins
+    SPV_AnyPtr:$pointer,
+    SPV_ScopeAttr:$memory_scope,
+    SPV_MemorySemanticsAttr:$semantics,
+    SPV_Float:$value
+  );
+
+  let results = (outs
+    SPV_Float:$result
+  );
+
+  let parser = [{ return ::parseAtomicUpdateOp(parser, result, true); }];
+  let printer = [{ return ::printAtomicUpdateOp(getOperation(), p); }];
+  let verifier = [{ return ::verifyAtomicUpdateOp<FloatType>(getOperation()); }];
+}
+
+// -----
+
 def SPV_AtomicIAddOp : SPV_AtomicUpdateWithValueOp<"AtomicIAdd", []> {
   let summary = [{
     Perform the following steps atomically with respect to any other atomic
     accesses within Scope to the same location:
   }];
 
   let description = [{
     1) load through Pointer to get an Original Value,

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

 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_OC_OpAtomicFAddEXT             : I32EnumAttrCase<"OpAtomicFAddEXT", 6035>;
 
 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_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_OpSubgroupBlockReadINTEL, SPV_OC_OpSubgroupBlockWriteINTEL
+      SPV_OC_OpSubgroupBlockReadINTEL, SPV_OC_OpSubgroupBlockWriteINTEL,
+      SPV_OC_OpAtomicFAddEXT
     ]>;
 
 // End opcode section. Generated from SPIR-V spec; DO NOT MODIFY!
 
 //===----------------------------------------------------------------------===//
 // SPIR-V op definitions
 //===----------------------------------------------------------------------===//
 

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

printer << spirv::stringifyScope(
           << "\" \"";
   auto memorySemanticsAttr = op->getAttrOfType<IntegerAttr>(kSemanticsAttrName);
   printer << spirv::stringifyMemorySemantics(
                  static_cast<spirv::MemorySemantics>(
                      memorySemanticsAttr.getInt()))
           << "\" " << op->getOperands() << " : " << op->getOperand(0).getType();
 }
 
+template <typename T>
+static StringRef stringifyTypeName();
+
+template <>
+StringRef stringifyTypeName<IntegerType>() {
+  return "integer";
+}
+
+template <>
+StringRef stringifyTypeName<FloatType>() {
+  return "float";
+}
+
 // Verifies an atomic update op.
+template <typename ExpectedElementType>
 static LogicalResult verifyAtomicUpdateOp(Operation *op) {
   auto ptrType = op->getOperand(0).getType().cast<spirv::PointerType>();
   auto elementType = ptrType.getPointeeType();
-  if (!elementType.isa<IntegerType>())
-    return op->emitOpError(
-               "pointer operand must point to an integer value, found ")
-           << elementType;
+  if (!elementType.isa<ExpectedElementType>())
+    return op->emitOpError() << "pointer operand must point to an "
+                             << stringifyTypeName<ExpectedElementType>()
+                             << " value, found " << elementType;
 
   if (op->getNumOperands() > 1) {
     auto valueType = op->getOperand(1).getType();
     if (valueType != elementType)
       return op->emitOpError("expected value to have the same type as the "
                              "pointer operand's pointee type ")
              << elementType << ", but found " << valueType;
   }

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

 // spv.AtomicXor
 //===----------------------------------------------------------------------===//
 
 func @atomic_xor(%ptr : !spv.ptr<i32, StorageBuffer>, %value : i32) -> i32 {
   // CHECK: spv.AtomicXor "Workgroup" "None" %{{.*}}, %{{.*}} : !spv.ptr<i32, StorageBuffer>
   %0 = spv.AtomicXor "Workgroup" "None" %ptr, %value : !spv.ptr<i32, StorageBuffer>
   return %0 : i32
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.AtomicFAddEXT
+//===----------------------------------------------------------------------===//
+
+func @atomic_fadd(%ptr : !spv.ptr<f32, StorageBuffer>, %value : f32) -> f32 {
+  // CHECK: spv.AtomicFAddEXT "Device" "None" %{{.*}}, %{{.*}} : !spv.ptr<f32, StorageBuffer>
+  %0 = spv.AtomicFAddEXT "Device" "None" %ptr, %value : !spv.ptr<f32, StorageBuffer>
+  return %0 : f32
+}
+
+// -----
+
+func @atomic_fadd(%ptr : !spv.ptr<i32, StorageBuffer>, %value : f32) -> f32 {
+  // expected-error @+1 {{pointer operand must point to an float value, found 'i32'}}
+  %0 = "spv.AtomicFAddEXT"(%ptr, %value) {memory_scope = 4: i32, semantics = 0x4 : i32} : (!spv.ptr<i32, StorageBuffer>, f32) -> (f32)
+  return %0 : f32
+}
+
+// -----
+
+func @atomic_fadd(%ptr : !spv.ptr<f32, StorageBuffer>, %value : f64) -> f64 {
+  // expected-error @+1 {{expected value to have the same type as the pointer operand's pointee type 'f32', but found 'f64'}}
+  %0 = "spv.AtomicFAddEXT"(%ptr, %value) {memory_scope = 2: i32, semantics = 0x8 : i32} : (!spv.ptr<f32, StorageBuffer>, f64) -> (f64)
+  return %0 : f64
+}
+
+// -----
+
+func @atomic_fadd(%ptr : !spv.ptr<f32, StorageBuffer>, %value : f32) -> f32 {
+  // expected-error @+1 {{expected at most one of these four memory constraints to be set: `Acquire`, `Release`,`AcquireRelease` or `SequentiallyConsistent`}}
+  %0 = spv.AtomicFAddEXT "Device" "Acquire|Release" %ptr, %value : !spv.ptr<f32, StorageBuffer>
+  return %0 : f32
+}

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

 // RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
-  // CHECK-LABEL: @atomic_compare_exchange_weak
-  spv.func @atomic_compare_exchange_weak(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 "None" {
+  // CHECK-LABEL: @test_int_atomics
+  spv.func @test_int_atomics(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 "None" {
     // CHECK: spv.AtomicCompareExchangeWeak "Workgroup" "Release" "Acquire" %{{.*}}, %{{.*}}, %{{.*}} : !spv.ptr<i32, Workgroup>
     %0 = spv.AtomicCompareExchangeWeak "Workgroup" "Release" "Acquire" %ptr, %value, %comparator: !spv.ptr<i32, Workgroup>
     // CHECK: spv.AtomicAnd "Device" "None" %{{.*}}, %{{.*}} : !spv.ptr<i32, Workgroup>
     %1 = spv.AtomicAnd "Device" "None" %ptr, %value : !spv.ptr<i32, Workgroup>
     // CHECK: spv.AtomicIAdd "Workgroup" "Acquire" %{{.*}}, %{{.*}} : !spv.ptr<i32, Workgroup>
     %2 = spv.AtomicIAdd "Workgroup" "Acquire" %ptr, %value : !spv.ptr<i32, Workgroup>
     // CHECK: spv.AtomicIDecrement "Workgroup" "Acquire" %{{.*}} : !spv.ptr<i32, Workgroup>
     %3 = spv.AtomicIDecrement "Workgroup" "Acquire" %ptr : !spv.ptr<i32, Workgroup>
spv.func @test_int_atomics(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 "None" {
     // CHECK: spv.AtomicXor "Workgroup" "AcquireRelease" %{{.*}}, %{{.*}} : !spv.ptr<i32, Workgroup>
     %11 = spv.AtomicXor "Workgroup" "AcquireRelease" %ptr, %value : !spv.ptr<i32, Workgroup>
     // CHECK: spv.AtomicCompareExchange "Workgroup" "Release" "Acquire" %{{.*}}, %{{.*}}, %{{.*}} : !spv.ptr<i32, Workgroup>
     %12 = spv.AtomicCompareExchange "Workgroup" "Release" "Acquire" %ptr, %value, %comparator: !spv.ptr<i32, Workgroup>
     // CHECK: spv.AtomicExchange "Workgroup" "Release" %{{.*}}, %{{.*}} : !spv.ptr<i32, Workgroup>
     %13 = spv.AtomicExchange "Workgroup" "Release" %ptr, %value: !spv.ptr<i32, Workgroup>
     spv.ReturnValue %0: i32
   }
+
+  // CHECK-LABEL: @test_float_atomics
+  spv.func @test_float_atomics(%ptr: !spv.ptr<f32, Workgroup>, %value: f32) -> f32 "None" {
+    // CHECK: spv.AtomicFAddEXT "Workgroup" "Acquire" %{{.*}}, %{{.*}} : !spv.ptr<f32, Workgroup>
+    %0 = spv.AtomicFAddEXT "Workgroup" "Acquire" %ptr, %value : !spv.ptr<f32, Workgroup>
+    spv.ReturnValue %0: f32
+  }
 }