[mlir][SPIR-V] Support floating-point atomic_rmw addf in MemRefToSPIRV

[aobolensk]

No description provided.

[llvmorg-github-actions]

@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Arseniy Obolenskiy (aobolensk)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/202330.diff

2 Files Affected:

• (modified) mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp (+7-8)
• (modified) mlir/test/Conversion/MemRefToSPIRV/atomic.mlir (+18)

diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index fe9f9c59e6ede..4674aa351315f 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
+++ b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
@@ -454,10 +454,6 @@ LogicalResult
 AtomicRMWOpPattern::matchAndRewrite(memref::AtomicRMWOp atomicOp,
                                     OpAdaptor adaptor,
                                     ConversionPatternRewriter &rewriter) const {
-  if (isa<FloatType>(atomicOp.getType()))
-    return rewriter.notifyMatchFailure(atomicOp,
-                                       "unimplemented floating-point case");
-
   auto memrefType = cast<MemRefType>(atomicOp.getMemref().getType());
   std::optional<spirv::Scope> scope = getAtomicOpScope(memrefType);
   if (!scope)
@@ -488,13 +484,13 @@ AtomicRMWOpPattern::matchAndRewrite(memref::AtomicRMWOp atomicOp,
                                        "failed to convert memref type");
 
   Type pointeeType = pointerType.getPointeeType();
-  auto dstType = dyn_cast<IntegerType>(
-      getElementTypeForStoragePointer(pointeeType, typeConverter));
-  if (!dstType)
+  Type storageElemType =
+      getElementTypeForStoragePointer(pointeeType, typeConverter);
+  if (!storageElemType || !storageElemType.isIntOrFloat())
     return rewriter.notifyMatchFailure(
         atomicOp, "failed to determine destination element type");
 
-  int dstBits = static_cast<int>(dstType.getWidth());
+  int dstBits = static_cast<int>(storageElemType.getIntOrFloatBitWidth());
   assert(dstBits % srcBits == 0);
 
   spirv::MemorySemantics memSem = getAtomicAcqRelMemorySemantics(memrefType);
@@ -509,6 +505,7 @@ AtomicRMWOpPattern::matchAndRewrite(memref::AtomicRMWOp atomicOp,
     break
 
     switch (atomicOp.getKind()) {
+      ATOMIC_CASE(addf, EXTAtomicFAddOp);
       ATOMIC_CASE(addi, AtomicIAddOp);
       ATOMIC_CASE(maxs, AtomicSMaxOp);
       ATOMIC_CASE(maxu, AtomicUMaxOp);
@@ -546,6 +543,8 @@ AtomicRMWOpPattern::matchAndRewrite(memref::AtomicRMWOp atomicOp,
         atomicOp,
         "sub-element-width atomic ops unsupported with Kernel capability");
 
+  auto dstType = cast<IntegerType>(storageElemType);
+
   auto accessChainOp = ptr.getDefiningOp<spirv::AccessChainOp>();
   if (!accessChainOp)
     return failure();
diff --git a/mlir/test/Conversion/MemRefToSPIRV/atomic.mlir b/mlir/test/Conversion/MemRefToSPIRV/atomic.mlir
index b5815a73ee8b2..fa416512aa144 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/atomic.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/atomic.mlir
@@ -146,3 +146,21 @@ func.func @atomic_andi_i8_storage_buffer(%value: i8, %memref: memref<16xi8, #spi
 }
 
 }
+
+// -----
+
+// Floating-point atomic add requires the shader_atomic_float_add extension.
+
+module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader, AtomicFloat32AddEXT], [SPV_EXT_shader_atomic_float_add]>, #spirv.resource_limits<>>} {
+
+//      CHECK: func.func @atomic_addf_storage_buffer
+// CHECK-SAME: (%[[VAL:.+]]: f32,
+func.func @atomic_addf_storage_buffer(%value: f32, %memref: memref<2x3x4xf32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> f32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.EXT.AtomicFAdd <Device> <AcquireRelease|UniformMemory> %[[AC]], %[[VAL]] : !spirv.ptr<f32, StorageBuffer>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "addf" %value, %memref[%i0, %i1, %i2] : (f32, memref<2x3x4xf32, #spirv.storage_class<StorageBuffer>>) -> f32
+  return %0: f32
+}
+
+}

[IgWod]

Why is it safe to cast float to int as the dstType?

[aobolensk]

Integers are handled above