[mlir][llvm] Changes to support amdgpu unsafe floating point atomic operations

Author: anchuraj

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -1899,6 +1899,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<BoolAttr>:$no_nans_fp_math,
     OptionalAttr<BoolAttr>:$approx_func_fp_math,
     OptionalAttr<BoolAttr>:$no_signed_zeros_fp_math,
+    OptionalAttr<BoolAttr>:$amdgpu_unsafe_fp_atomics,
     OptionalAttr<StrAttr>:$denormal_fp_math,
     OptionalAttr<StrAttr>:$denormal_fp_math_f32,
     OptionalAttr<StrAttr>:$fp_contract,

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -2140,6 +2140,7 @@ static constexpr std::array kExplicitAttributes{
     StringLiteral("aarch64_pstate_sm_body"),
     StringLiteral("aarch64_pstate_sm_compatible"),
     StringLiteral("aarch64_pstate_sm_enabled"),
+    StringLiteral("amdgpu-unsafe-fp-atomics"),
     StringLiteral("alwaysinline"),
     StringLiteral("approx-func-fp-math"),
     StringLiteral("convergent"),
@@ -2256,6 +2257,8 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
   else if (func->hasFnAttribute("aarch64_preserves_za"))
     funcOp.setArmPreservesZa(true);
 
+  if (func->hasFnAttribute("amdgpu-unsafe-fp-atomics"))
+    funcOp.setAmdgpuUnsafeFpAtomics(true);
   llvm::Attribute attr = func->getFnAttribute(llvm::Attribute::VScaleRange);
   if (attr.isValid()) {
     MLIRContext *context = funcOp.getContext();

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -1496,6 +1496,10 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   else if (func.getArmPreservesZa())
     llvmFunc->addFnAttr("aarch64_preserves_za");
 
+  if (auto amdgpuUnsafeFpAtomics = func.getAmdgpuUnsafeFpAtomics())
+    llvmFunc->addFnAttr("amdgpu-unsafe-fp-atomics",
+                        llvm::toStringRef(*amdgpuUnsafeFpAtomics));
+
   if (auto targetCpu = func.getTargetCpu())
     llvmFunc->addFnAttr("target-cpu", *targetCpu);
 

mlir/test/Dialect/LLVMIR/func.mlir

@@ -236,6 +236,12 @@ module {
     llvm.return
   }
 
+  llvm.func @amdgpu_fp_unsafe_atomics() attributes {amdgpu_fp_unsafe_atomics = true} {
+    // CHECK: @amdgpu_fp_unsafe_atomics
+    // CHECK-SAME: attributes {amdgpu_fp_unsafe_atomics = true}
+    llvm.return
+  }
+
   // CHECK: llvm.comdat @__llvm_comdat
   llvm.comdat @__llvm_comdat {
     // CHECK: llvm.comdat_selector @any any

mlir/test/Target/LLVMIR/Import/function-attributes.ll

@@ -267,6 +267,12 @@ define void @arm_preserves_za_func() "aarch64_preserves_za" {
   ret void
 }
 
+; // -----
+
+; CHECK-LABEL: @func_attr_amdgpu_unsafe_fp_atomics_true
+; CHECK-SAME: attributes {amdgpu_unsafe_fp_atomics = true}
+declare void @func_attr_amdgpu_unsafe_fp_atomics_true() "amdgpu-unsafe-fp-atomics"="true"
+
 // -----
 
 ; CHECK-LABEL: @section_func

mlir/test/Target/LLVMIR/llvmir.mlir

@@ -2474,6 +2474,15 @@ llvm.func @preserves_za_func() attributes {arm_preserves_za} {
 
 // -----
 
+// CHECK-LABEL: @amdgpu_unsafe_fp_atomics_func
+// CHECK-SAME: #[[ATTR:[0-9]*]]
+llvm.func @amdgpu_unsafe_fp_atomics_func() attributes {amdgpu_unsafe_fp_atomics = true} {
+  llvm.return
+}
+// CHECK: #[[ATTR]] = { "amdgpu-unsafe-fp-atomics"="true" }
+
+// -----
+
 //
 // frame pointer attribute.
 //