feat(rocBLAS): default to hipBLASLt for ZGEMM/CGEMM/SGEMM/fp16/bf16 on MI300A (gfx942, 228 CU)

projects/rocblas/docs/conceptual/rocblas-design-notes.rst

@@ -49,7 +49,8 @@ the data type, and the problem. The architectures which default to hipBLASLt, an
         default to Tensile.
 
     * - gfx942 (MI300A, 228 CUs)
-      - ``f64_r`` only. ``f16_r``, ``bf16_r``, and all other data types default to Tensile.
+      - ``f16_r``, ``bf16_r``, ``f64_r``, ``f32_r``, ``f32_c``, and ``f64_c``, including their mixed-precision
+        ``gemm_ex`` forms. All other data types default to Tensile.
 
     * - All others
       - None. Tensile is always the default.

projects/rocblas/library/src/hipblaslt_host.cpp

@@ -86,6 +86,20 @@ namespace
     template <>
     constexpr auto hipblaslt_compute_type<double> = HIPBLAS_COMPUTE_64F;
 
+    // Compute type honouring the handle's math mode; mirrors the Tensile path's
+    // setF32XdlMathOp(XFloat32). Without it the hipBLASLt backend silently runs
+    // full fp32 while rocblas_get_math_mode still reports xf32.
+    template <typename Ti, typename Tc>
+    hipblasComputeType_t hipblaslt_compute_type_for(rocblas_handle handle)
+    {
+        if constexpr(std::is_same<Ti, float>{} && std::is_same<Tc, float>{})
+        {
+            if(handle && handle->math_mode == rocblas_xf32_xdl_math_op)
+                return HIPBLAS_COMPUTE_32F_FAST_TF32;
+        }
+        return hipblaslt_compute_type<Tc>;
+    }
+
     /**************************************************************************
     * We normally print error messages only once, to avoid excessive logging *
     **************************************************************************/
@@ -251,7 +265,7 @@ namespace
                                  hipblaslt_datatype<Ti>,
                                  hipblaslt_datatype<To>,
                                  hipblaslt_datatype<To>,
-                                 hipblaslt_compute_type<Tc>);
+                                 hipblaslt_compute_type_for<Ti, Tc>(prob.handle));
 
         hipblaslt_ext::GemmProblemType problemType;
         problemType.setOpA((hipblasOperation_t)prob.trans_a);
@@ -260,7 +274,7 @@ namespace
         problemType.setTypeB(hipblaslt_datatype<Ti>);
         problemType.setTypeC(hipblaslt_datatype<To>);
         problemType.setTypeD(hipblaslt_datatype<To>);
-        problemType.setTypeCompute(hipblaslt_compute_type<Tc>);
+        problemType.setTypeCompute(hipblaslt_compute_type_for<Ti, Tc>(prob.handle));
 
         hipblaslt_ext::GemmEpilogue epilogue;
         hipblaslt_ext::GemmInputs   inputs;
@@ -304,7 +318,7 @@ namespace
                                         hipblaslt_datatype<Ti>,
                                         hipblaslt_datatype<To>,
                                         hipblaslt_datatype<To>,
-                                        hipblaslt_compute_type<Tc>);
+                                        hipblaslt_compute_type_for<Ti, Tc>(prob.handle));
 
         try
         {
@@ -315,7 +329,7 @@ namespace
             problemType.setTypeB(hipblaslt_datatype<Ti>);
             problemType.setTypeC(hipblaslt_datatype<To>);
             problemType.setTypeD(hipblaslt_datatype<To>);
-            problemType.setTypeCompute(hipblaslt_compute_type<Tc>);
+            problemType.setTypeCompute(hipblaslt_compute_type_for<Ti, Tc>(prob.handle));
 
             std::vector<int64_t>                     Ms(prob.batch_count);
             std::vector<int64_t>                     Ns(prob.batch_count);
@@ -695,7 +709,7 @@ rocblas_status runContractionProblemHipBlasLT(const RocblasContractionProblem<Ti
             }
         }
         THROW_IF_HIPBLASLT_ERROR(hipblasLtMatmulDescCreate(
-            &matmulDesc, hipblaslt_compute_type<Tc>, hipblaslt_scaletype<Ti>));
+            &matmulDesc, hipblaslt_compute_type_for<Ti, Tc>(prob.handle), hipblaslt_scaletype<Ti>));
         THROW_IF_HIPBLASLT_ERROR(hipblasLtMatmulDescSetAttribute(
             matmulDesc, HIPBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(int32_t)));
         THROW_IF_HIPBLASLT_ERROR(hipblasLtMatmulDescSetAttribute(
@@ -1018,7 +1032,7 @@ rocblas_status getAllSolutionsHipBlasLT(const RocblasContractionProblem<Ti, To,
                                                             hipblaslt_datatype<Ti>,
                                                             hipblaslt_datatype<To>,
                                                             hipblaslt_datatype<To>,
-                                                            hipblaslt_compute_type<Tc>,
+                                                            hipblaslt_compute_type_for<Ti, Tc>(prob.handle),
                                                             heuristicResults_temp);
 
                     heuristicResults.insert(heuristicResults.end(),
@@ -1083,7 +1097,7 @@ rocblas_status getAllSolutionsHipBlasLT(const RocblasContractionProblem<Ti, To,
                                                     hipblaslt_datatype<Ti>,
                                                     hipblaslt_datatype<To>,
                                                     hipblaslt_datatype<To>,
-                                                    hipblaslt_compute_type<Tc>,
+                                                    hipblaslt_compute_type_for<Ti, Tc>(prob.handle),
                                                     heuristicResults);
 
             std::shared_ptr<hipblaslt_ext::GemmInstance> gemm;

projects/rocblas/library/src/tensile_host.cpp

@@ -1300,17 +1300,24 @@ bool useHipBLASLt(const RocblasContractionProblem<Ti, To, Tc>& prob)
                 return false;
         }
 
-        // gfx942: hipBLASLt is used for fp64 on all gfx942 devices, and for fp16/bf16
-        // only on MI300X (304 CUs); MI300A (228 CUs) stays on Tensile for fp16/bf16.
-        // TODO remove after all types are supported
+        // gfx942: the set of types defaulting to hipBLASLt depends on the device.
+        //   every gfx942 (MI300X and MI300A): fp16/bf16/fp64
+        //   228-CU MI300A only:               complex-fp64, fp32, complex-fp32
+        // TODO expand once more gfx942 configs are validated
         if(arch == 942)
         {
-            if constexpr(sizeof(Ti) == 2)
+            constexpr bool any_gfx942 = sizeof(Ti) == 2 || std::is_same<Ti, double>::value;
+
+            constexpr bool only_MI300A = std::is_same<Ti, rocblas_double_complex>::value
+                                         || std::is_same<Ti, float>::value
+                                         || std::is_same<Ti, rocblas_float_complex>::value;
+
+            if constexpr(only_MI300A)
             {
-                if(prob.handle->device_properties.multiProcessorCount != 304)
+                if(prob.handle->device_properties.multiProcessorCount != 228)
                     return false;
             }
-            else if constexpr(!std::is_same<Ti, double>::value)
+            else if constexpr(!any_gfx942)
             {
                 return false;
             }