[mlir] Remove convertible identity restriction for memref.atomic_rmw …

[Max191]

…to LLVM

[llvmbot]

@llvm/pr-subscribers-mlir

Author: None (Max191)

Changes

…to LLVM

---

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

2 Files Affected:

• (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (-2)
• (modified) mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir (+18)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 91b1210efec23e0..5d97cb8425947f2 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -1562,8 +1562,6 @@ struct AtomicRMWOpLowering : public LoadStoreOpLowering<memref::AtomicRMWOp> {
   LogicalResult
   matchAndRewrite(memref::AtomicRMWOp atomicOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (failed(match(atomicOp)))
-      return failure();
     auto maybeKind = matchSimpleAtomicOp(atomicOp);
     if (!maybeKind)
       return failure();
diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
index 3b3a51d609be972..37999d6fc14ad19 100644
--- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
@@ -400,6 +400,24 @@ func.func @atomic_rmw(%I : memref<10xi32>, %ival : i32, %F : memref<10xf32>, %fv
 
 // -----
 
+func.func @atomic_rmw_with_offset(%I : memref<10xi32, strided<[1], offset: 5>>, %ival : i32, %i : index) {
+  memref.atomic_rmw andi %ival, %I[%i] : (i32, memref<10xi32, strided<[1], offset: 5>>) -> i32
+  return
+}
+// CHECK-LABEL:  func @atomic_rmw_with_offset
+// CHECK-SAME:   %[[ARG0:.+]]: memref<10xi32, strided<[1], offset: 5>>
+// CHECK-SAME:   %[[ARG1:.+]]: i32
+// CHECK-SAME:   %[[ARG2:.+]]: index
+// CHECK:        %[[MEMREF_STRUCT:.+]] = builtin.unrealized_conversion_cast %[[ARG0]] : memref<10xi32, strided<[1], offset: 5>> to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+// CHECK:        %[[INDEX:.+]] = builtin.unrealized_conversion_cast %[[ARG2]] : index to i64
+// CHECK:        %[[BASE_PTR:.+]] = llvm.extractvalue %[[MEMREF_STRUCT]][1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
+// CHECK:        %[[OFFSET:.+]] = llvm.mlir.constant(5 : index) : i64
+// CHECK:        %[[OFFSET_PTR:.+]] = llvm.getelementptr %[[BASE_PTR]][%[[OFFSET]]] : (!llvm.ptr, i64) -> !llvm.ptr, i32
+// CHECK:        %[[PTR:.+]] = llvm.getelementptr %[[OFFSET_PTR]][%[[INDEX]]] : (!llvm.ptr, i64) -> !llvm.ptr, i32
+// CHECK:        llvm.atomicrmw _and %[[PTR]], %[[ARG1]] acq_rel
+
+// -----
+
 // CHECK-LABEL: func @generic_atomic_rmw
 func.func @generic_atomic_rmw(%I : memref<10xi32>, %i : index) {
   %x = memref.generic_atomic_rmw %I[%i] : memref<10xi32> {

[Max191]

memref.atomic_rmw will fail for memref types that have an offset because they do not have identity maps. However, I don't see why this restriction is needed, and this restriction causes some cases to fail to lower to LLVM. memref.store has no such restriction, and the conversion for store is the same as atomic_rmw.

Unless there was a reason for having the restriction there, it would be good to remove it to unblock some lowerings, especially once #72181 lands and the narrow type emulation of memref.store starts introducing memref.atomic_rmw

[MaheshRavishankar]

I actually dont understand this check..... This looks OK to me, but what was stopped by this change?

[hanhanW]

It looks okay to me. I don't know who should be involved in the review even I looked at git blame. It looks like very very few people update the pattern and op...

[Max191]

I actually dont understand this check..... This looks OK to me, but what was stopped by this change?

It fails when the memref type has an offset because it checks that the memref type has an identity map. But the pattern can handle offsets, so it should be fine to remove the check

[MaheshRavishankar]

It seems like this only works for strided memref types. Can you check that getStridesAndOffset on the source type returns success, i.e. return failure when that method fails. I looked into the implementation of getStridedElementPtr and it asserts if thats not the case, so letting that through would lead to an assert

[MaheshRavishankar]

It seems like this only works for strided memref types. Can you check that getStridesAndOffset on the source type returns success, i.e. return failure when that method fails. I looked into the implementation of getStridedElementPtr and it asserts if thats not the case, so letting that through would lead to an assert