[SYCL][JM] Add support for FP8 and FP4 for CRI

clang/lib/CodeGen/CodeGenTypes.cpp

@@ -399,6 +399,11 @@ llvm::Type *CodeGenTypes::ConvertSPVCooperativeMatrixType(RecordDecl *RD) {
       CompTy = llvm::Type::getFloatTy(getLLVMContext());
     } else if (LlvmTyName == "bfloat16") {
       CompTy = llvm::Type::getInt16Ty(getLLVMContext());
+    } else if (LlvmTyName == "fp8_e5m2_x" || LlvmTyName == "fp8_e4m3_x" ||
+               LlvmTyName == "fp8_e8m0_x") {
+      CompTy = llvm::Type::getInt8Ty(getLLVMContext());
+    } else if (LlvmTyName == "fp4_e2m1_x") {
+      CompTy = llvm::Type::getIntNTy(getLLVMContext(), 4);
     } else {
       llvm_unreachable("Wrong matrix base type!");
     }

sycl/include/sycl/__spirv/spirv_ops_matrix.hpp

@@ -151,4 +151,67 @@ extern __DPCPP_SYCL_EXTERNAL void __spirv_CooperativeMatrixPrefetchINTEL(
     T *Ptr, uint32_t NumRows, uint32_t NumCols, unsigned int CacheLevel,
     __spv::MatrixLayout Layout, size_t Stride);
 
+// FP4E2M1 Upconversion
+template <typename To, typename From, std::size_t R, std::size_t C,
+          __spv::MatrixUse U,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern __DPCPP_SYCL_EXTERNAL
+    __spv::__spirv_CooperativeMatrixKHR<To, S, R, C, U> *
+    __spirv_ConvertFP4E2M1ToHF16INTEL(
+        __spv::__spirv_CooperativeMatrixKHR<From, S, R, C, U> *Object);
+template <typename To, typename From, std::size_t R, std::size_t C,
+          __spv::MatrixUse U,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern __DPCPP_SYCL_EXTERNAL
+    __spv::__spirv_CooperativeMatrixKHR<To, S, R, C, U> *
+    __spirv_ConvertFP4E2M1ToBF16INTEL(
+        __spv::__spirv_CooperativeMatrixKHR<From, S, R, C, U> *Object);
+template <typename To, typename From, std::size_t R, std::size_t C,
+          __spv::MatrixUse U,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern __DPCPP_SYCL_EXTERNAL
+    __spv::__spirv_CooperativeMatrixKHR<To, S, R, C, U> *
+    __spirv_ConvertFP4E2M1ToHF8INTEL(
+        __spv::__spirv_CooperativeMatrixKHR<From, S, R, C, U> *Object);
+template <typename To, typename From, std::size_t R, std::size_t C,
+          __spv::MatrixUse U,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern __DPCPP_SYCL_EXTERNAL
+    __spv::__spirv_CooperativeMatrixKHR<To, S, R, C, U> *
+    __spirv_ConvertFP4E2M1ToBF8INTEL(
+        __spv::__spirv_CooperativeMatrixKHR<From, S, R, C, U> *Object);
+// FP4E2M1 down conversion
+template <typename To, typename From, std::size_t R, std::size_t C,
+          __spv::MatrixUse U,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern __DPCPP_SYCL_EXTERNAL
+    __spv::__spirv_CooperativeMatrixKHR<To, S, R, C, U> *
+    __spirv_ConvertHF16ToFP4E2M1INTEL(
+        __spv::__spirv_CooperativeMatrixKHR<From, S, R, C, U> *Object);
+template <typename To, typename From, std::size_t R, std::size_t C,
+          __spv::MatrixUse U,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern __DPCPP_SYCL_EXTERNAL
+    __spv::__spirv_CooperativeMatrixKHR<To, S, R, C, U> *
+    __spirv_ConvertBF16ToFP4E2M1INTEL(
+        __spv::__spirv_CooperativeMatrixKHR<From, S, R, C, U> *Object);
+
+template <typename TA, typename TB, typename TC, typename TAS, typename TBS,
+          std::size_t M, std::size_t K, std::size_t N, __spv::MatrixUse UA,
+          __spv::MatrixUse UB, __spv::MatrixUse UC, __spv::MatrixUse UAS,
+          __spv::MatrixUse UBS,
+          __spv::MatrixLayout LA = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LB = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LC = __spv::MatrixLayout::RowMajor,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern __DPCPP_SYCL_EXTERNAL
+    __spv::__spirv_CooperativeMatrixKHR<TC, S, M, N, UC> *
+    __spirv_CooperativeMatrixMulAddScaledINTEL(
+        __spv::__spirv_CooperativeMatrixKHR<TA, S, M, K, UA> *A,
+        __spv::__spirv_CooperativeMatrixKHR<TB, S, K, N, UB> *B,
+        __spv::__spirv_CooperativeMatrixKHR<TC, S, M, N, UC> *C,
+        __spv::__spirv_CooperativeMatrixKHR<TAS, S, M, K / 32, UAS> *Ascale,
+        __spv::__spirv_CooperativeMatrixKHR<TBS, S, K / 32, N, UBS> *Bscale,
+        size_t Operands = 0);
+
 #endif

sycl/include/sycl/__spirv/spirv_types.hpp

@@ -130,7 +130,13 @@ enum class MatrixOperands : uint32_t {
   MatrixAAndBTF32ComponentsINTEL = 0x20,
   MatrixAAndBBFloat16ComponentsINTEL = 0x40,
   MatrixCBFloat16ComponentsINTEL = 0x80,
-  MatrixResultBFloat16ComponentsINTEL = 0x100
+  MatrixResultBFloat16ComponentsINTEL = 0x100,
+  MatrixABFloat8ComponentsINTEL = 0x400,
+  MatrixBBFloat8ComponentsINTEL = 0x800,
+  MatrixAHFloat8ComponentsINTEL = 0x1000,
+  MatrixBHFloat8ComponentsINTEL = 0x2000,
+  MatrixAFP4S1E2M1ComponentsINTEL = 0x10000,
+  MatrixBFP4S1E2M1ComponentsINTEL = 0x20000
 };
 
 template <typename T, Scope::Flag S = Scope::Flag::Subgroup, std::size_t R = 1,

sycl/include/sycl/ext/oneapi/matrix/matrix-unified-utils.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <cstddef>                      // size_t
 #include <optional>                     // std::optional
 #include <string_view>                  // std::string_view
 #include <sycl/__spirv/spirv_types.hpp> // __spv namespace
@@ -19,9 +20,21 @@ inline namespace _V1 {
 namespace ext {
 namespace oneapi {
 namespace experimental {
+
+// The matrix interfaces only ever name the 4-bit and 8-bit floating point types
+// as template arguments, so forward declarations are enough here. Including
+// <sycl/ext/oneapi/experimental/float_4bit/types.hpp> and its 8-bit counterpart
+// instead would pull both extensions into the transitive closure of
+// <sycl/sycl.hpp>, which they are deliberately not part of: code that uses
+// fp4/fp8 includes those headers explicitly.
+template <size_t N> class fp4_e2m1_x;
+template <size_t N> class fp8_e4m3_x;
+template <size_t N> class fp8_e5m2_x;
+template <size_t N> class fp8_e8m0_x;
+
 namespace matrix {
 
-enum class use { a, b, accumulator };
+enum class use { a, b, accumulator, scale };
 
 enum class layout {
   row_major = 0,
@@ -49,6 +62,7 @@ constexpr UseToUseStringPair UseToUseStringMap[] = {
     {ext::oneapi::experimental::matrix::use::a, "use::a"},
     {ext::oneapi::experimental::matrix::use::b, "use::b"},
     {ext::oneapi::experimental::matrix::use::accumulator, "use::accumulator"},
+    {ext::oneapi::experimental::matrix::use::scale, "use::scale"},
 };
 
 constexpr const char *
@@ -85,6 +99,11 @@ extern "C" constexpr __spv::MatrixLayout joint_matrix_layout_to_spv(
   }
 }
 
+template <typename T> struct is_fp4_e2m1 : std::false_type {};
+template <size_t N>
+struct is_fp4_e2m1<sycl::ext::oneapi::experimental::fp4_e2m1_x<N>>
+    : std::true_type {};
+
 template <typename Ta, typename Tb, typename Tc, typename Td>
 constexpr uint32_t CalculateMatrixOperand() {
   uint32_t returnValue = 0x00;
@@ -104,6 +123,44 @@ constexpr uint32_t CalculateMatrixOperand() {
   if constexpr (std::is_signed<Tb>::value)
     returnValue += static_cast<uint32_t>(
         __spv::MatrixOperands::MatrixBSignedComponentsKHR);
+  if constexpr (
+      std::is_same<Ta, sycl::ext::oneapi::experimental::fp8_e5m2_x<1>>::value &&
+      std::is_same<Tb, sycl::ext::oneapi::experimental::fp8_e5m2_x<1>>::value &&
+      std::is_same<Tc, float>::value)
+    returnValue += static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixABFloat8ComponentsINTEL) +
+                   static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixBBFloat8ComponentsINTEL);
+  if constexpr (
+      std::is_same<Ta, sycl::ext::oneapi::experimental::fp8_e5m2_x<1>>::value &&
+      std::is_same<Tb, sycl::ext::oneapi::experimental::fp8_e4m3_x<1>>::value &&
+      std::is_same<Tc, float>::value)
+    returnValue += static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixABFloat8ComponentsINTEL) +
+                   static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixBHFloat8ComponentsINTEL);
+  if constexpr (
+      std::is_same<Ta, sycl::ext::oneapi::experimental::fp8_e4m3_x<1>>::value &&
+      std::is_same<Tb, sycl::ext::oneapi::experimental::fp8_e5m2_x<1>>::value &&
+      std::is_same<Tc, float>::value)
+    returnValue += static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixAHFloat8ComponentsINTEL) +
+                   static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixBBFloat8ComponentsINTEL);
+  if constexpr (
+      std::is_same<Ta, sycl::ext::oneapi::experimental::fp8_e4m3_x<1>>::value &&
+      std::is_same<Tb, sycl::ext::oneapi::experimental::fp8_e4m3_x<1>>::value &&
+      std::is_same<Tc, float>::value)
+    returnValue += static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixAHFloat8ComponentsINTEL) +
+                   static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixBHFloat8ComponentsINTEL);
+  if constexpr (sycl::detail::is_fp4_e2m1<Ta>::value &&
+                sycl::detail::is_fp4_e2m1<Tb>::value)
+    returnValue += static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixAFP4S1E2M1ComponentsINTEL) +
+                   static_cast<uint32_t>(
+                       __spv::MatrixOperands::MatrixBFP4S1E2M1ComponentsINTEL);
   return returnValue;
 }
 

sycl/include/sycl/ext/oneapi/matrix/matrix-unified.hpp

@@ -576,6 +576,75 @@ joint_matrix_prefetch(Group sg, T *Ptr, size_t stride,
 #endif // defined(__SYCL_DEVICE_ONLY__)
 }
 
+template <typename To, typename From, typename Group, std::size_t M,
+          std::size_t N, use Use, layout Layout>
+inline __SYCL_ALWAYS_INLINE void
+joint_matrix_convert(Group,
+                     const joint_matrix<Group, From, Use, M, N, Layout> &src,
+                     joint_matrix<Group, To, Use, M, N, Layout> &dst) {
+#if defined(__SYCL_DEVICE_ONLY__)
+  // FP4E2M1 Upconversion
+  if constexpr (sycl::detail::is_fp4_e2m1<From>::value) {
+    if constexpr (std::is_same<To, sycl::half>::value)
+      dst.spvm = __spirv_ConvertFP4E2M1ToHF16INTEL<To>(src.spvm);
+    else if constexpr (std::is_same<To, sycl::ext::oneapi::bfloat16>::value)
+      dst.spvm = __spirv_ConvertFP4E2M1ToBF16INTEL<To>(src.spvm);
+    else if constexpr (std::is_same<To, sycl::ext::oneapi::experimental::
+                                            fp8_e4m3_x<1>>::value)
+      dst.spvm = __spirv_ConvertFP4E2M1ToHF8INTEL<To>(src.spvm);
+    else if constexpr (std::is_same<To, sycl::ext::oneapi::experimental::
+                                            fp8_e5m2_x<1>>::value)
+      dst.spvm = __spirv_ConvertFP4E2M1ToBF8INTEL<To>(src.spvm);
+  }
+  // FP4E2M1 down conversion
+  else if constexpr (sycl::detail::is_fp4_e2m1<To>::value) {
+    if constexpr (std::is_same<From, sycl::half>::value)
+      dst.spvm = __spirv_ConvertHF16ToFP4E2M1INTEL<To>(src.spvm);
+    else if constexpr (std::is_same<From, sycl::ext::oneapi::bfloat16>::value)
+      dst.spvm = __spirv_ConvertBF16ToFP4E2M1INTEL<To>(src.spvm);
+  }
+#else
+  std::ignore = src;
+  std::ignore = dst;
+  std::ignore = Layout;
+  throw exception(make_error_code(errc::runtime),
+                  "joint_matrix_convert is not supported on host.");
+#endif // defined(__SYCL_DEVICE_ONLY__)
+}
+
+template <typename Group, typename Ta, typename Tb, typename Tc, typename Td,
+          std::size_t M, std::size_t K, std::size_t N, layout LayoutA,
+          layout LayoutB, layout LayoutAs, layout LayoutBs>
+inline __SYCL_ALWAYS_INLINE void joint_matrix_bmad(
+    Group,
+    joint_matrix<Group, Td, use::accumulator, M, N,
+                 sycl::ext::oneapi::experimental::matrix::layout::dynamic> &D,
+    const joint_matrix<Group, Ta, use::a, M, K, LayoutA> &A,
+    const joint_matrix<Group, Tb, use::b, K, N, LayoutB> &B,
+    const joint_matrix<Group, sycl::ext::oneapi::experimental::fp8_e8m0_x<1>,
+                       use::scale, M, K / 32, LayoutAs> &Ascale,
+    const joint_matrix<Group, sycl::ext::oneapi::experimental::fp8_e8m0_x<1>,
+                       use::scale, K / 32, N, LayoutBs> &Bscale,
+    const joint_matrix<Group, Tc, use::accumulator, M, N,
+                       sycl::ext::oneapi::experimental::matrix::layout::dynamic>
+        &C) {
+#if defined(__SYCL_DEVICE_ONLY__)
+  constexpr uint32_t MatrixOperand =
+      sycl::detail::CalculateMatrixOperand<Ta, Tb, Tc, Td>();
+  D.spvm = __spirv_CooperativeMatrixMulAddScaledINTEL(
+      A.spvm, B.spvm, C.spvm, Ascale.spvm, Bscale.spvm, MatrixOperand);
+#else
+  std::ignore = A;
+  std::ignore = B;
+  std::ignore = Ascale;
+  std::ignore = Bscale;
+  std::ignore = C;
+  std::ignore = D;
+  throw exception(make_error_code(errc::runtime),
+                  "joint matrix is not supported on host.");
+#endif // defined(__SYCL_DEVICE_ONLY__)
+}
+
 } // namespace matrix
 } // namespace experimental
 } // namespace oneapi

sycl/include/sycl/ext/oneapi/matrix/query-types.hpp

@@ -31,7 +31,10 @@ enum class matrix_type {
   uint8,
   uint16,
   uint32,
-  uint64
+  uint64,
+  fp8_e5m2,
+  fp8_e4m3,
+  fp4_e2m1
 };
 
 struct combination {
@@ -64,7 +67,12 @@ struct matrix_combinations
 // Type to matrix type string conversion used in compile-time
 namespace detail {
 template <typename T> constexpr const char *convertTypeToMatrixTypeString() {
-  return "";
+  // fp4_e2m1_x<N> is a class template, so it cannot be handled by an explicit
+  // function template specialization; match on the element format instead.
+  if constexpr (is_fp4_e2m1<T>::value)
+    return "matrix_type::fp4_e2m1";
+  else
+    return "";
 }
 template <>
 constexpr const char *
@@ -79,6 +87,18 @@ constexpr const char *convertTypeToMatrixTypeString<
     sycl::ext::oneapi::experimental::matrix::precision::tf32>() {
   return "matrix_type::tf32";
 }
+template <>
+constexpr const char *convertTypeToMatrixTypeString<
+    sycl::ext::oneapi::experimental::fp8_e5m2_x<1>>() {
+  return "matrix_type::fp8_e5m2";
+}
+
+template <>
+constexpr const char *convertTypeToMatrixTypeString<
+    sycl::ext::oneapi::experimental::fp8_e4m3_x<1>>() {
+  return "matrix_type::fp8_e4m3";
+}
+
 template <> constexpr const char *convertTypeToMatrixTypeString<float>() {
   return "matrix_type::fp32";
 }

sycl/source/detail/device_impl.hpp

@@ -2206,6 +2206,33 @@ class device_impl {
           {8, 0, 0, 0, 16, 8, matrix_type::tf32, matrix_type::tf32,
            matrix_type::fp32, matrix_type::fp32},
       };
+      // fp8 and fp4_e2m1 are supported starting with Crescent Island.
+      bool SupportsFP8AndFP4 = (architecture::intel_gpu_cri == DeviceArch);
+      if (SupportsFP8AndFP4)
+        pvc_combs.insert(
+            pvc_combs.end(),
+            {
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e5m2,
+                 matrix_type::fp8_e5m2, matrix_type::bf16, matrix_type::bf16},
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e5m2,
+                 matrix_type::fp8_e4m3, matrix_type::bf16, matrix_type::bf16},
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e4m3,
+                 matrix_type::fp8_e5m2, matrix_type::bf16, matrix_type::bf16},
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e4m3,
+                 matrix_type::fp8_e4m3, matrix_type::bf16, matrix_type::bf16},
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e5m2,
+                 matrix_type::fp8_e5m2, matrix_type::fp32, matrix_type::fp32},
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e5m2,
+                 matrix_type::fp8_e4m3, matrix_type::fp32, matrix_type::fp32},
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e4m3,
+                 matrix_type::fp8_e5m2, matrix_type::fp32, matrix_type::fp32},
+                {8, 0, 0, 0, 16, 32, matrix_type::fp8_e4m3,
+                 matrix_type::fp8_e4m3, matrix_type::fp32, matrix_type::fp32},
+                {0, 0, 0, 8, 16, 32, matrix_type::fp4_e2m1,
+                 matrix_type::fp4_e2m1, matrix_type::fp32, matrix_type::fp32},
+                {0, 0, 0, 8, 16, 32, matrix_type::fp4_e2m1,
+                 matrix_type::fp4_e2m1, matrix_type::bf16, matrix_type::bf16},
+            });
       return pvc_combs;
     } else if ((architecture::intel_gpu_dg2_g10 == DeviceArch) ||
                (architecture::intel_gpu_dg2_g11 == DeviceArch) ||

sycl/source/detail/program_manager/program_manager.cpp

@@ -3105,6 +3105,12 @@ convertMatrixTypeStringMatrixTypeEnumValue(
     return matrix_ext::matrix_type::uint32;
   else if ("uint64" == MatrixTypeStringView)
     return matrix_ext::matrix_type::uint64;
+  else if ("fp8_e5m2" == MatrixTypeStringView)
+    return matrix_ext::matrix_type::fp8_e5m2;
+  else if ("fp8_e4m3" == MatrixTypeStringView)
+    return matrix_ext::matrix_type::fp8_e4m3;
+  else if ("fp4_e2m1" == MatrixTypeStringView)
+    return matrix_ext::matrix_type::fp4_e2m1;
   return std::nullopt;
 }
 
@@ -3188,6 +3194,9 @@ std::optional<sycl::exception> checkDevSupportJointMatrix(
             Combination.nsize);
         break;
       }
+      case matrix_ext::use::scale:
+        // TODO as part of a new query
+        break;
       }
 
       // early exit if we have a match