Force FP8 gemms into F16 dot

Author: phambinhfin

xla/service/gpu/transforms/gemm_rewriter.cc

@@ -1152,6 +1152,33 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
                   << PrimitiveType_Name(b_type);
           return false;
         }
+        if (a_type == F8E4M3FN && b_type == F8E4M3FN) {
+          VLOG(1) << "Failed to rewrite " << instr->ToShortString()
+                  << " into FP8 Custom Call. For "
+                  << rocm_compute_capability.gfx_version()
+                  << " arch, one of the input types must be F8E4M3FN, but got "
+                  << PrimitiveType_Name(a_type) << " and "
+                  << PrimitiveType_Name(b_type);
+          return false;
+        }
+        if (a_type == F8E5M2 && b_type == F8E4M3FN) {
+          VLOG(1) << "Failed to rewrite " << instr->ToShortString()
+                  << " into FP8 Custom Call. For "
+                  << rocm_compute_capability.gfx_version()
+                  << " arch, one of the input types must be F8E4M3FN, but got "
+                  << PrimitiveType_Name(a_type) << " and "
+                  << PrimitiveType_Name(b_type);
+          return false;
+        }
+                if (a_type == F8E4M3FN && b_type == F8E4M3FN) {
+          VLOG(1) << "Failed to rewrite " << instr->ToShortString()
+                  << " into FP8 Custom Call. For "
+                  << rocm_compute_capability.gfx_version()
+                  << " arch, one of the input types must be F8E4M3FN, but got "
+                  << PrimitiveType_Name(a_type) << " and "
+                  << PrimitiveType_Name(b_type);
+          return false;
+        }
         if ((a_type != F8E5M2 && a_type != F8E4M3FN) ||
             (b_type != F8E5M2 && b_type != F8E4M3FN)) {
           VLOG(1)
@@ -2632,4 +2659,4 @@ absl::StatusOr<bool> GemmRewriter::Run(
 }
 
 }  // namespace gpu
-}  // namespace xla
+}  // namespace xla