[mlir][spirv] Support i32 memref.atomic_rmw conversion

[antiagainst]

Summary

These cases can be directly mapped to spirv.AtomicI* ops.

[kuhar]

LGTM. Sorry for the belated review, I missed the notification.

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 //===- MemRefToSPIRV.cpp - MemRef to SPIR-V Patterns ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements patterns to convert MemRef dialect to SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 #define DEBUG_TYPE "memref-to-spirv-pattern"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//

 public:
   using OpConversionPattern<memref::AllocOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::AllocOp operation, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts memref.automic_rmw operations to SPIR-V atomic operations.
+class AtomicRMWOpPattern final
+    : public OpConversionPattern<memref::AtomicRMWOp> {
+public:
+  using OpConversionPattern<memref::AtomicRMWOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::AtomicRMWOp atomicOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Removed a deallocation if it is a supported allocation. Currently only
 /// removes deallocation if the memory space is workgroup memory.
 class DeallocOpPattern final : public OpConversionPattern<memref::DeallocOp> {
 public:
   using OpConversionPattern<memref::DeallocOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::DeallocOp operation, OpAdaptor adaptor,
AllocOpPattern::matchAndRewrite(memref::AllocOp operation, OpAdaptor adaptor,
   }
 
   // Get pointer to global variable at the current scope.
   rewriter.replaceOpWithNewOp<spirv::AddressOfOp>(operation, varOp);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// AllocOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+AtomicRMWOpPattern::matchAndRewrite(memref::AtomicRMWOp atomicOp,
+                                    OpAdaptor adaptor,
+                                    ConversionPatternRewriter &rewriter) const {
+  if (atomicOp.getType().isa<FloatType>())
+    return rewriter.notifyMatchFailure(atomicOp,
+                                       "unimplemented floating-point case");
+
+  auto memrefType = atomicOp.getMemref().getType().cast<MemRefType>();
+  std::optional<spirv::Scope> scope = getAtomicOpScope(memrefType);
+  if (!scope)
+    return rewriter.notifyMatchFailure(atomicOp,
+                                       "unsupported memref memory space");
+
+  auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
+  Type resultType = typeConverter.convertType(atomicOp.getType());
+  if (!resultType)
+    return rewriter.notifyMatchFailure(atomicOp,
+                                       "failed to convert result type");
+
+  auto loc = atomicOp.getLoc();
+  Value ptr =
+      spirv::getElementPtr(typeConverter, memrefType, adaptor.getMemref(),
+                           adaptor.getIndices(), loc, rewriter);
+
+  if (!ptr)
+    return failure();
+
+#define ATOMIC_CASE(kind, spirvOp)                                             \
+  case arith::AtomicRMWKind::kind:                                             \
+    rewriter.replaceOpWithNewOp<spirv::spirvOp>(                               \
+        atomicOp, resultType, ptr, *scope,                                     \
+        spirv::MemorySemantics::AcquireRelease, adaptor.getValue());           \
+    break
+
+  switch (atomicOp.getKind()) {
+    ATOMIC_CASE(addi, AtomicIAddOp);
+    ATOMIC_CASE(maxs, AtomicSMaxOp);
+    ATOMIC_CASE(maxu, AtomicUMaxOp);
+    ATOMIC_CASE(mins, AtomicSMinOp);
+    ATOMIC_CASE(minu, AtomicUMinOp);
+    ATOMIC_CASE(ori, AtomicOrOp);
+    ATOMIC_CASE(andi, AtomicAndOp);
+  default:
+    return rewriter.notifyMatchFailure(atomicOp, "unimplemented atomic kind");
+  }
+
+#undef ATOMIC_CASE
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 DeallocOpPattern::matchAndRewrite(memref::DeallocOp operation,
                                   OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
   MemRefType deallocType = operation.getMemref().getType().cast<MemRefType>();

 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 void populateMemRefToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
-  patterns.add<AllocaOpPattern, AllocOpPattern, DeallocOpPattern,
-               IntLoadOpPattern, IntStoreOpPattern, LoadOpPattern,
-               MemorySpaceCastOpPattern, StoreOpPattern>(typeConverter,
-                                                         patterns.getContext());
+  patterns.add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
+               DeallocOpPattern, IntLoadOpPattern, IntStoreOpPattern,
+               LoadOpPattern, MemorySpaceCastOpPattern, StoreOpPattern>(
+      typeConverter, patterns.getContext());
 }
 } // namespace mlir

mlir/test/Conversion/MemRefToSPIRV/atomic.mlir

+// RUN: mlir-opt -split-input-file -convert-memref-to-spirv %s -o - | FileCheck %s
+
+module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader], []>, #spirv.resource_limits<>>} {
+
+//      CHECK: func.func @atomic_addi_storage_buffer
+// CHECK-SAME: (%[[VAL:.+]]: i32,
+func.func @atomic_addi_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicIAdd "Device" "AcquireRelease" %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "addi" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
+  return %0: i32
+}
+
+//      CHECK: func.func @atomic_maxs_workgroup
+// CHECK-SAME: (%[[VAL:.+]]: i32,
+func.func @atomic_maxs_workgroup(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<Workgroup>>, %i0: index, %i1: index, %i2: index) -> i32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicSMax "Workgroup" "AcquireRelease" %[[AC]], %[[VAL]] : !spirv.ptr<i32, Workgroup>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "maxs" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<Workgroup>>) -> i32
+  return %0: i32
+}
+
+//      CHECK: func.func @atomic_maxu_storage_buffer
+// CHECK-SAME: (%[[VAL:.+]]: i32,
+func.func @atomic_maxu_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicUMax "Device" "AcquireRelease" %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "maxu" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
+  return %0: i32
+}
+
+//      CHECK: func.func @atomic_mins_workgroup
+// CHECK-SAME: (%[[VAL:.+]]: i32,
+func.func @atomic_mins_workgroup(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<Workgroup>>, %i0: index, %i1: index, %i2: index) -> i32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicSMin "Workgroup" "AcquireRelease" %[[AC]], %[[VAL]] : !spirv.ptr<i32, Workgroup>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "mins" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<Workgroup>>) -> i32
+  return %0: i32
+}
+
+//      CHECK: func.func @atomic_minu_storage_buffer
+// CHECK-SAME: (%[[VAL:.+]]: i32,
+func.func @atomic_minu_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicUMin "Device" "AcquireRelease" %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "minu" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
+  return %0: i32
+}
+
+//      CHECK: func.func @atomic_ori_workgroup
+// CHECK-SAME: (%[[VAL:.+]]: i32,
+func.func @atomic_ori_workgroup(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<Workgroup>>, %i0: index, %i1: index, %i2: index) -> i32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicOr "Workgroup" "AcquireRelease" %[[AC]], %[[VAL]] : !spirv.ptr<i32, Workgroup>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "ori" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<Workgroup>>) -> i32
+  return %0: i32
+}
+
+//      CHECK: func.func @atomic_andi_storage_buffer
+// CHECK-SAME: (%[[VAL:.+]]: i32,
+func.func @atomic_andi_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
+  // CHECK: %[[AC:.+]] = spirv.AccessChain
+  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicAnd "Device" "AcquireRelease" %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
+  // CHECK: return %[[ATOMIC]]
+  %0 = memref.atomic_rmw "andi" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
+  return %0: i32
+}
+
+}
+