[Mlir-commits] [mlir] [mlir][MemRefToLLVM] Support floating-point types in GenericAtomicRMWOp lowering (PR #194300)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Fangcao Wang (LittleMeepo)

<details>
<summary>Changes</summary>

`llvm.cmpxchg` only accepts integer or pointer operands. When the memref element type is floating-point (e.g. `f16`), bitcast values to a same-width integer for the CAS and bitcast the new-loaded result back to the original float type.

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


2 Files Affected:

- (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+28-5) 
- (modified) mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir (+31) 


``````````diff
diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index c42a85fa375ba..61069fc4d660d 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -711,10 +711,20 @@ struct GenericAtomicRMWOpLowering
     auto loc = atomicOp.getLoc();
     Type valueType = typeConverter->convertType(atomicOp.getResult().getType());
 
+    // `llvm.cmpxchg` only supports integer or pointer operands. For
+    // floating-point element types, perform the CAS on a same-width integer
+    // and bitcast at the boundaries.
+    bool needsBitcast = isa<FloatType>(valueType);
+    Type cmpxchgType = valueType;
+    if (needsBitcast) {
+      unsigned bitWidth = cast<FloatType>(valueType).getWidth();
+      cmpxchgType = rewriter.getIntegerType(bitWidth);
+    }
+
     // Split the block into initial, loop, and ending parts.
     auto *initBlock = rewriter.getInsertionBlock();
     auto *loopBlock = rewriter.splitBlock(initBlock, Block::iterator(atomicOp));
-    loopBlock->addArgument(valueType, loc);
+    loopBlock->addArgument(cmpxchgType, loc);
 
     auto *endBlock =
         rewriter.splitBlock(loopBlock, Block::iterator(atomicOp)++);
@@ -727,15 +737,21 @@ struct GenericAtomicRMWOpLowering
     Value init = LLVM::LoadOp::create(
         rewriter, loc, typeConverter->convertType(memRefType.getElementType()),
         dataPtr);
+    if (needsBitcast)
+      init = LLVM::BitcastOp::create(rewriter, loc, cmpxchgType, init);
     LLVM::BrOp::create(rewriter, loc, init, loopBlock);
 
     // Prepare the body of the loop block.
     rewriter.setInsertionPointToStart(loopBlock);
 
     // Clone the GenericAtomicRMWOp region and extract the result.
-    auto loopArgument = loopBlock->getArgument(0);
+    Value loopArgument = loopBlock->getArgument(0);
+    Value loopArgForBody = loopArgument;
+    if (needsBitcast)
+      loopArgForBody =
+          LLVM::BitcastOp::create(rewriter, loc, valueType, loopArgument);
     IRMapping mapping;
-    mapping.map(atomicOp.getCurrentValue(), loopArgument);
+    mapping.map(atomicOp.getCurrentValue(), loopArgForBody);
     Block &entryBlock = atomicOp.body().front();
     for (auto &nestedOp : entryBlock.without_terminator()) {
       Operation *clone = rewriter.clone(nestedOp, mapping);
@@ -747,6 +763,8 @@ struct GenericAtomicRMWOpLowering
     if (!result) {
       return atomicOp.emitError("result not defined in region");
     }
+    if (needsBitcast)
+      result = LLVM::BitcastOp::create(rewriter, loc, cmpxchgType, result);
 
     // Prepare the epilog of the loop block.
     // Append the cmpxchg op to the end of the loop block.
@@ -763,9 +781,14 @@ struct GenericAtomicRMWOpLowering
     LLVM::CondBrOp::create(rewriter, loc, ok, endBlock, ArrayRef<Value>(),
                            loopBlock, newLoaded);
 
+    // The 'result' of the atomic_rmw op is the newly loaded value. Bitcast
+    // back to the float type if needed. Insert at the start of `endBlock` so
+    // the bitcast precedes the existing terminator (split into endBlock).
+    if (needsBitcast) {
+      rewriter.setInsertionPointToStart(endBlock);
+      newLoaded = LLVM::BitcastOp::create(rewriter, loc, valueType, newLoaded);
+    }
     rewriter.setInsertionPointToEnd(endBlock);
-
-    // The 'result' of the atomic_rmw op is the newly loaded value.
     rewriter.replaceOp(atomicOp, {newLoaded});
 
     return success();
diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
index d2fe5ab582b71..1c0025e6cd406 100644
--- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
@@ -547,6 +547,37 @@ llvm.func @generic_atomic_rmw() {
 
 // -----
 
+// CHECK-LABEL: func @generic_atomic_rmw_f16
+// CHECK-INTERFACE-LABEL: func @generic_atomic_rmw_f16
+llvm.func @generic_atomic_rmw_f16() {
+  %I = "test.foo"() : () -> (memref<10xf16>)
+  %i = "test.foo"() : () -> (index)
+  %x = memref.generic_atomic_rmw %I[%i] : memref<10xf16> {
+    ^bb0(%old_value : f16):
+      memref.atomic_yield %old_value : f16
+  }
+  llvm.return
+}
+// `llvm.cmpxchg` only accepts integer/pointer operands, so floating-point
+// values are bitcast to a same-width integer for the CAS and bitcast back.
+// CHECK:        %[[INIT:.*]] = llvm.load %{{.*}} : !llvm.ptr -> f16
+// CHECK-NEXT:   %[[INIT_INT:.*]] = llvm.bitcast %[[INIT]] : f16 to i16
+// CHECK-NEXT:   llvm.br ^bb1(%[[INIT_INT]] : i16)
+// CHECK-NEXT: ^bb1(%[[LOADED:.*]]: i16):
+// CHECK-NEXT:   %[[LOADED_F16:.*]] = llvm.bitcast %[[LOADED]] : i16 to f16
+// CHECK-NEXT:   %[[RES_INT:.*]] = llvm.bitcast %[[LOADED_F16]] : f16 to i16
+// CHECK-NEXT:   %[[PAIR:.*]] = llvm.cmpxchg %{{.*}}, %[[LOADED]], %[[RES_INT]]
+// CHECK-SAME:                      acq_rel monotonic : !llvm.ptr, i16
+// CHECK-NEXT:   %[[NEW:.*]] = llvm.extractvalue %[[PAIR]][0]
+// CHECK-NEXT:   %[[OK:.*]] = llvm.extractvalue %[[PAIR]][1]
+// CHECK-NEXT:   llvm.cond_br %[[OK]], ^bb2, ^bb1(%[[NEW]] : i16)
+// CHECK-NEXT: ^bb2:
+// CHECK-NEXT:   %{{.*}} = llvm.bitcast %[[NEW]] : i16 to f16
+
+// CHECK-INTERFACE: llvm.cmpxchg
+
+// -----
+
 // CHECK-LABEL: func @generic_atomic_rmw_in_alloca_scope
 // CHECK-INTERFACE-LABEL: func @generic_atomic_rmw_in_alloca_scope
 llvm.func @generic_atomic_rmw_in_alloca_scope() {

``````````

</details>


https://github.com/llvm/llvm-project/pull/194300