Have fp8 GEMMs go to hipblaslt only (#3990)

Will always use hipblaslt for fp8 GEMMs with this PR.
FP8 will be eliminated if hipblaslt is not avaliable.

Author: CharlieL7

src/targets/gpu/device_name.cpp

@@ -87,6 +87,20 @@ bool gfx_default_rocblas()
 }
 #endif
 
+bool hipblaslt_supported()
+{
+#if !MIGRAPHX_USE_HIPBLASLT
+    return false;
+#else
+    const auto device_name = trim(split_string(get_device_name(), ':').front());
+    // hipblaslt is supported for MI200 and above, and Navi3x and above.
+    return (device_name == "gfx90a" or
+            (starts_with(device_name, "gfx94") and device_name >= "gfx942") or
+            (starts_with(device_name, "gfx95") and device_name >= "gfx950") or
+            starts_with(device_name, "gfx110") or starts_with(device_name, "gfx120"));
+#endif
+}
+
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/targets/gpu/hipblaslt.cpp

@@ -53,16 +53,6 @@ hipblaslt_preference_ptr create_hipblaslt_preference_ptr()
     return hipblaslt_preference_ptr{preference};
 }
 
-bool hipblaslt_supported()
-{
-    const auto device_name = trim(split_string(get_device_name(), ':').front());
-    // hipblaslt is supported for MI200 and above, and Navi3x and above.
-    return (device_name == "gfx90a" or
-            (starts_with(device_name, "gfx94") and device_name >= "gfx942") or
-            (starts_with(device_name, "gfx95") and device_name >= "gfx950") or
-            starts_with(device_name, "gfx110") or starts_with(device_name, "gfx120"));
-}
-
 #endif // MIGRAPHX_USE_HIPBLASLT
 
 } // namespace gpu

src/targets/gpu/include/migraphx/gpu/device_name.hpp

@@ -47,6 +47,8 @@ MIGRAPHX_GPU_EXPORT bool gfx_has_fp8fnuz_support();
 
 MIGRAPHX_GPU_EXPORT bool gfx_default_rocblas();
 
+MIGRAPHX_GPU_EXPORT bool hipblaslt_supported();
+
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/targets/gpu/include/migraphx/gpu/hipblaslt.hpp

@@ -100,7 +100,6 @@ using hipblaslt_preference_ptr = MIGRAPHX_MANAGE_PTR(hipblasLtMatmulPreference_t
 
 hipblaslt_handle_ptr create_hipblaslt_handle_ptr();
 hipblaslt_preference_ptr create_hipblaslt_preference_ptr();
-MIGRAPHX_GPU_EXPORT bool hipblaslt_supported();
 const size_t hipblaslt_workspace_size = 2 * 128 * 1024 * 1024;
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS

src/targets/gpu/lowering.cpp

@@ -33,6 +33,7 @@
 #include <migraphx/pass_manager.hpp>
 #include <migraphx/iterator_for.hpp>
 #include <migraphx/program.hpp>
+#include <migraphx/fp8_types.hpp>
 
 #include <migraphx/op/common.hpp>
 #include <migraphx/op/dot.hpp>
@@ -253,17 +254,20 @@ struct miopen_apply
             auto output = insert_allocation(ins, ins->get_shape());
             refs.push_back(output);
 
-#if MIGRAPHX_USE_HIPBLASLT
+            bool has_fp8_inputs =
+                std::any_of(ins->inputs().begin(), ins->inputs().end(), [](auto i_input) {
+                    return contains(fp8_types{}.get(), i_input->get_shape().type());
+                });
+
             // Check if user explicitly sets rocBLAS as GEMM provider, or
             // if the hardware cannot support hipblaslt, or
             // if the hardware is defaulted to use rocBLAS (such as gfx90).
-            if((string_value_of(MIGRAPHX_SET_GEMM_PROVIDER{}) == "rocblas") or
-               not hipblaslt_supported() or gpu::gfx_default_rocblas())
+            if(not has_fp8_inputs and
+               ((string_value_of(MIGRAPHX_SET_GEMM_PROVIDER{}) == "rocblas") or
+                not hipblaslt_supported() or gpu::gfx_default_rocblas()))
             {
-#endif
                 return mod->replace_instruction(
                     ins, rocblas_gemm<Op>{Op{}, 1, 0, compute_fp32}, refs);
-#if MIGRAPHX_USE_HIPBLASLT
             }
             std::string op_name = "gpu::hip_gemm";
             if(contains(name, "quant_"))
@@ -277,7 +281,6 @@ struct miopen_apply
                 ins->inputs().at(0),
                 ins->inputs().at(1),
                 output);
-#endif
         });
     }
 #endif

src/targets/gpu/target.cpp

@@ -112,13 +112,15 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
     }
 
     // whiltelist supported Ops for the FP8 types
-    // rocBLAS does not support any FP8 types
     std::set<std::string> unsupported_fp8fnuz_ops = {};
-    if(string_value_of(MIGRAPHX_SET_GEMM_PROVIDER{}) == "rocblas" or gpu::gfx_default_rocblas())
+
+    // disable dot & quant_dot if no hipblaslt
+    if(not hipblaslt_supported())
     {
         unsupported_fp8fnuz_ops.insert("dot");
         unsupported_fp8fnuz_ops.insert("quant_dot");
     }
+
 #if MIGRAPHX_USE_MIOPEN // MIOpen doesn't have support for fp8 pooling yet.
     unsupported_fp8fnuz_ops.insert("pooling");
 #endif
@@ -141,6 +143,14 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
     unsupported_fp8fnuz_ops.insert("argmin");
 
     std::set<std::string> unsupported_fp8ocp_ops = {};
+
+    // disable dot & quant_dot if no hipblaslt
+    if(not hipblaslt_supported())
+    {
+        unsupported_fp8ocp_ops.insert("dot");
+        unsupported_fp8ocp_ops.insert("quant_dot");
+    }
+
 #if MIGRAPHX_USE_MIOPEN
     // MIOpen doesn't have support for fp8 pooling yet.
     unsupported_fp8ocp_ops.insert("pooling");