use hip_has_flags

ROCm · Sep 12, 2023 · a64ee16 · a64ee16
1 parent c109d1e
commit a64ee16
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/src/targets/gpu/CMakeLists.txt b/src/targets/gpu/CMakeLists.txt
@@ -204,12 +204,6 @@ if(MIGRAPHX_ENABLE_MLIR)
     target_link_libraries(migraphx_gpu PRIVATE rocMLIR::rockCompiler)
 endif()
 
-math(EXPR hip_VERSION_FLAT "(${hip_VERSION_MAJOR} * 1000 + ${hip_VERSION_MINOR}) * 100000 + ${hip_VERSION_PATCH}")
-if(${hip_VERSION_FLAT} GREATER_EQUAL 500723302)
-    set(MIGRAPHX_HIP_REQUIRE_UNIFORM_WG ON CACHE BOOL "Whether HIP requires uniform workgroup sizes or not")
-    target_compile_definitions(migraphx_gpu PUBLIC -DMIGRAPHX_HIP_REQUIRE_UNIFORM_WG)
-endif()
-
 if(MIGRAPHX_USE_HIPRTC)
     message(STATUS "MIGraphX is using hipRTC")
     target_compile_definitions(migraphx_gpu PRIVATE -DMIGRAPHX_USE_HIPRTC=1)

diff --git a/src/targets/gpu/compile_hip_code_object.cpp b/src/targets/gpu/compile_hip_code_object.cpp
@@ -139,17 +139,28 @@ void hip_compile_options::set_launch_params(
         global = compute_global(local);
 }
 
+static bool hip_accept_non_uniform_wg()
+{
+    static bool non_uniform_wg = hip_has_flags({"-fno-offload-uniform-block"});
+    return non_uniform_wg;
+}
+
 std::function<std::size_t(std::size_t local)>
 compute_global_for(context& ctx, std::size_t n, std::size_t over)
 {
     assert(over > 0);
     std::size_t max_global = ctx.get_current_device().get_cu_count() *
                              ctx.get_current_device().get_max_workitems_per_cu();
     return [n, over, max_global](std::size_t local) {
-        std::size_t groups     = 1 + (n - 1) / local;
+        std::size_t num_elements = n;
+        if(not hip_accept_non_uniform_wg())
+        {
+            num_elements = (1 + (n - 1) / local) * local;
+        }
+        std::size_t groups     = 1 + (num_elements - 1) / local;
         std::size_t max_blocks = max_global / local;
         std::size_t nglobal    = std::min(max_blocks * over, groups) * local;
-        return std::min(nglobal, n);
+        return std::min(nglobal, num_elements);
     };
 }
 
@@ -183,12 +194,10 @@ operation compile_hip_code_object(const std::string& content, hip_compile_option
         generate_args_hpp(options.virtual_inputs.empty() ? options.inputs : options.virtual_inputs);
     srcs.push_back(src_file{fs::path{"args.hpp"},
                             std::make_pair(args_hpp.data(), args_hpp.data() + args_hpp.size())});
-#ifdef MIGRAPHX_HIP_REQUIRE_UNIFORM_WG
-    if(options.global % options.local != 0)
-    {
+    if(options.global % options.local != 0 and hip_accept_non_uniform_wg())
         options.params += " -fno-offload-uniform-block";
-    }
-#endif
+    else
+        assert(options.global % options.local == 0);
     options.params += " -DMIGRAPHX_NGLOBAL=" + std::to_string(options.global);
     options.params += " -DMIGRAPHX_NLOCAL=" + std::to_string(options.local);
     options.params += " " + join_strings(compiler_warnings(), " ");