[AMD] Add an environment variable to disable buffer atomics (#7311)

Adds an environment variable, `AMDGCN_DISABLE_BUFFER_ATOMICS`
that can be set to disable Buffer atomics on the MI300. We have seen
significant wins from BufferOps in general, but we haven't benchmarked
the impact of atomics at Meta and we have seen some issue internally.

We are working to resolve these issues/will open reproducers upstream,
but with BufferOps being the default we would be much happier if we
could allow a production workload to opt to accept BufferOps but not
atomics.

Author: njriasan

include/triton/Tools/Sys/GetEnv.hpp

@@ -14,6 +14,7 @@ namespace mlir::triton {
 inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     // clang-format off
     "AMDGCN_ENABLE_DUMP",
+    "AMDGCN_USE_BUFFER_ATOMICS",
     "AMDGCN_USE_BUFFER_OPS",
     "DISABLE_LLVM_OPT",
     "DISABLE_MMA_V3",

python/triton/knobs.py

@@ -436,6 +436,8 @@ class nvidia_knobs(base_knobs):
 
 class amd_knobs(base_knobs):
     use_buffer_ops: env_bool = env_bool("AMDGCN_USE_BUFFER_OPS", True)
+    # Note: This requires use_buffer_ops be true to have any effect
+    use_buffer_atomics: env_bool = env_bool("AMDGCN_USE_BUFFER_ATOMICS", True)
     dump_amdgcn: env_bool = env_bool("AMDGCN_ENABLE_DUMP")
     libhip_path: env_opt_str = env_opt_str("TRITON_LIBHIP_PATH")
     lld_path: env_opt_str = env_opt_str("TRITON_HIP_LLD_PATH")

third_party/amd/backend/compiler.py

@@ -261,7 +261,7 @@ def make_ttgir(mod, metadata, options):
         if knobs.amd.use_buffer_ops:
             amd.passes.ttgpuir.add_canonicalize_pointers(pm)
             passes.common.add_canonicalizer(pm)
-            amd.passes.ttgpuir.add_convert_to_buffer_ops(pm, options.arch)
+            amd.passes.ttgpuir.add_convert_to_buffer_ops(pm, options.arch, knobs.amd.use_buffer_atomics)
 
         amd.passes.ttgpuir.add_fold_true_cmpi(pm)
         passes.common.add_canonicalizer(pm)

third_party/amd/include/TritonAMDGPUTransforms/Passes.td

@@ -137,6 +137,9 @@ def TritonAMDGPUConvertToBufferOps : Pass<"tritonamdgpu-convert-buffer-ops", "ml
     Option<"archGenerationName", "arch-generation-name",
            "std::string", /*default=*/"std::string{}",
            "GFX generation name of target device.">,
+    Option<"allowBufferAtomics", "allow-buffer-atomics",
+           "bool", /*default*/"true",
+           "Allow buffer atomic operations when the hardware supports it.">,
   ];
 }
 

third_party/amd/lib/TritonAMDGPUTransforms/ConvertToBufferOps.cpp

@@ -557,7 +557,7 @@ class TritonAMDGPUConvertToBufferOpsPass
     // lowering to LLVM
     triton::AMD::ISAFamily isaFamily =
         triton::AMD::deduceISAFamily(archGenerationName);
-    if (ISAFamily::CDNA3 == isaFamily)
+    if (this->allowBufferAtomics && ISAFamily::CDNA3 == isaFamily)
       patterns.add<ConvertTritonAtomicRMWOpToBufferAtomicRMW>(
           context, assumptions, axisInfoAnalysis, solver, isaFamily);
 

third_party/amd/python/triton_amd.cc

@@ -69,9 +69,9 @@ void init_triton_amd_passes_ttgpuir(py::module &&m) {
     pm.addNestedPass<mlir::triton::FuncOp>(
         mlir::createTritonAMDGPUCanonicalizePointers());
   });
-  ADD_PASS_OPTION_WRAPPER_1("add_convert_to_buffer_ops",
+  ADD_PASS_OPTION_WRAPPER_2("add_convert_to_buffer_ops",
                             mlir::createTritonAMDGPUConvertToBufferOps,
-                            const std::string &);
+                            const std::string &, bool);
   ADD_PASS_WRAPPER_0("add_reorder_instructions",
                      mlir::createTritonAMDGPUReorderInstructions);
   ADD_PASS_WRAPPER_0("add_fold_true_cmpi", mlir::createTritonAMDFoldTrueCmpI);