[GPU]Xetla support MTL (#176)

* bugfix

* add AOT

* Update examples/CMakeLists.txt

Co-authored-by: Yi DING <[email protected]>

* XETLA_PRINTF replace cout

* SLMSIZE 128KB 64KB

* Update include/subgroup/tile/impl/load_xe.hpp

Co-authored-by: Meng, Hengyu <[email protected]>

* use arch_attr_t

* add more shape for int4

* bugfix dump_mat

* save

---------

Co-authored-by: Yi DING <[email protected]>
Co-authored-by: Meng, Hengyu <[email protected]>

Author: 3 people

CMakeLists.txt

@@ -43,9 +43,16 @@ if (${LOG} STREQUAL "on")
     add_definitions(-DLOG_PRINT)
 endif ()
 
+# AOT device
+set(AOT_DEVICE "" CACHE STRING "Set device list for AOT build")
+
 add_compile_options(-fsycl)
 add_link_options(-fsycl)
 if(UNIX)
+    if (AOT_DEVICE)
+        add_compile_options(-fsycl-targets=spir64_gen)
+        add_link_options(-fsycl-targets=spir64_gen -Xs "-device ${AOT_DEVICE}") # MTL
+    endif()
     add_compile_options(-fp-model=precise -Wall -Wextra -Werror)
     add_link_options(-lmkl_intel_lp64 -lmkl_sequential -lmkl_core -lpthread -lm)
     link_libraries(-lgtest -lgtest_main)

examples/CMakeLists.txt

@@ -2,7 +2,7 @@ include_directories(${CMAKE_SOURCE_DIR}/include)
 include_directories(${CMAKE_SOURCE_DIR})
 
 # Creates a separate device code module for each SYCL* kernel
-# so that kernel for Dg2 and Xe will be JIT separately
+# so that kernel for XeHpc, XeHpg, and XeLpg will be JIT separately
 add_compile_options(-fsycl-device-code-split=per_kernel)
 add_link_options(-fsycl-device-code-split=per_kernel)
 

include/common/core/arch_config.hpp

@@ -108,10 +108,10 @@ struct register_attr_t {};
 template <grf_mode grf_num_mode, gpu_arch arch_tag>
 struct client_register_attr_base_t {
   static constexpr uint32_t acc_reg_in_bytes =
-      (grf_num_mode == grf_mode::normal) ? 4 * 32 : 8 * 32;
+      (grf_num_mode == grf_mode::normal) ? 4 * 64 : 8 * 64;
   static constexpr uint32_t grf_in_bytes =
-      (grf_num_mode == grf_mode::normal) ? 128 * 32 : 256 * 32;
-  static constexpr uint32_t reg_in_bytes = 32;
+      (grf_num_mode == grf_mode::normal) ? 128 * 64 : 256 * 64;
+  static constexpr uint32_t reg_in_bytes = 64;
 };
 
 template <grf_mode grf_num_mode>
@@ -139,7 +139,7 @@ struct client_arch_attr_base_t {
   template <msg_type message_type = msg_type::block_2d>
   using load_store_attr = load_store_attr_t<message_type, gpu_arch::XeHpg>;
 
-  template <grf_mode grf_num_mode = grf_mode::normal>
+  template <grf_mode grf_num_mode = grf_mode::double_grf>
   using register_attr = register_attr_t<grf_num_mode, gpu_arch::XeHpg>;
 
   using mma_attr = mma_attr_t<gpu_arch::XeHpg>;

include/experimental/group/gemm/compute_policy.hpp

@@ -82,8 +82,9 @@ struct compute_policy_int4_dequantize<
 
   static constexpr bool is_int4_matB_policy = true;
 
-  static constexpr uint32_t block_size_x_b =
-      arch_attr_t<arch_tag>::mma_attr::mma_n_in_elem;
+  static constexpr uint32_t block_size_x_b = (mma_engine == mma_engine::xmx)
+      ? arch_attr_t<arch_tag>::mma_attr::mma_n_in_elem
+      : 32;
   static constexpr uint32_t block_bytes_y_b = 32;
   static_assert(
       block_bytes_x_a == block_bytes_y_b,

include/experimental/group/gemm/impl/int4_dequantize_xe.hpp

@@ -560,20 +560,6 @@ class gemm_t<
   }
 
  private:
-  template <typename T>
-  void dump_mat(T mat, size_t tile_x, size_t tile_y) {
-#pragma unroll
-    for (size_t row = 0; row < tile_x; row++) {
-#pragma unroll
-      for (size_t col = 0; col < tile_y; col++) {
-        sycl::ext::oneapi::experimental::printf(
-            "%0.1f ", (float)(sycl::half)mat.reg[row * tile_y + col]);
-      }
-      sycl::ext::oneapi::experimental::printf("\n ");
-    }
-    sycl::ext::oneapi::experimental::printf("\n ");
-  }
-
   inline void dequantize(
       matB_acc_t& matB_acc,
       matB_t& matB,

include/experimental/kernel/gemm/impl/int4_dequantize_kslicing_xe.hpp

@@ -451,9 +451,11 @@ class gemm_universal_t<
   static cl::sycl::range<3> get_local_range() {
     uint32_t local_range_m = (wg_tile_m + sg_tile_m - 1) / sg_tile_m;
     uint32_t local_range_n = (wg_tile_n + sg_tile_n - 1) / sg_tile_n;
-    // std::cout << "Local range: {" << num_local_kslicing << ", " <<
-    // local_range_m
-    //           << ", " << local_range_n << "} \n";
+    XETLA_PRINTF(
+        "Local range: {%d, %d, %d}",
+        num_local_kslicing,
+        local_range_m,
+        local_range_n);
     assert(local_range_m * local_range_n * num_local_kslicing <= 32);
     return cl::sycl::range<3>{num_local_kslicing, local_range_m, local_range_n};
   };
@@ -471,8 +473,11 @@ class gemm_universal_t<
     uint32_t group_range_m = (matrix_m + wg_tile_m - 1) / wg_tile_m;
     uint32_t group_range_n = (matrix_n + wg_tile_n - 1) / wg_tile_n;
     group_swizzle_t::update_group_range(group_range_m, group_range_n);
-    // std::cout << "Group range: {" << num_global_kslicing << ", "
-    //           << group_range_m << ", " << group_range_n << "} \n";
+    XETLA_PRINTF(
+        "Group range: {%d, %d, %d}",
+        num_global_kslicing,
+        group_range_m,
+        group_range_n);
     return cl::sycl::range<3>{
         num_global_kslicing, group_range_m, group_range_n};
   };

include/subgroup/tile/impl/load_xe.hpp

@@ -28,6 +28,7 @@ namespace gpu::xetla::subgroup {
 namespace detail {
 template <typename tile_t, typename payload_t>
 struct check_load_type {
+  static constexpr bool is_lsc_gather = true;
   static constexpr bool is_global_block_2d =
       (payload_t::memory_space == mem_space::global &&
        (payload_t::message_type == msg_type::block_2d) &&
@@ -444,6 +445,7 @@ template <
     typename payload_t>
 __XETLA_API typename std::enable_if_t<
     detail::check_load_type<tile_t, payload_t>::is_global_block_2d &&
+    detail::check_load_type<tile_t, payload_t>::is_lsc_gather &&
     payload_t::arch_tag <= gpu_arch::XeHpg>
 tile_load(tile_t& tile, payload_t& payload) {
   using dtype = typename payload_t::dtype;
@@ -531,6 +533,77 @@ tile_load(tile_t& tile, payload_t& payload) {
   }
 }
 
+/// @brief This function loads data from unaligned-2D memory surface.
+/// Loads an array of rectangular regions (X,Y)..(X+W,Y+H) from memory into
+/// registers. Each block will be loaded serially by its corresponding payload.
+/// @tparam tile_t Is the tile_t struct contains registers.
+/// These registers will be the destination of load operation.
+/// @tparam payload_t Is the mem_payload_t struct describing the memory
+/// information. Payload indicates the source of load operation.
+/// @tparam L1 Is the cache hint for L1 cache.
+/// @tparam L3 Is the cache hint for L3 cache.
+/// @param tile Is the tile object with type tile_t, holds the return data of
+/// the loads.
+/// @param payload Is the payload object with type payload_t. Contains all the
+/// information for loads.
+/// @return No return, update in place.
+template <
+    cache_hint L1 = cache_hint::cached,
+    cache_hint L3 = cache_hint::cached,
+    typename tile_t,
+    typename payload_t>
+__XETLA_API typename std::enable_if_t<
+    detail::check_load_type<tile_t, payload_t>::is_global_block_2d &&
+    !detail::check_load_type<tile_t, payload_t>::is_lsc_gather &&
+    !arch_has_2d_load_store(payload_t::arch_tag)>
+tile_load(tile_t& tile, payload_t& payload) {
+  using dtype = typename payload_t::dtype;
+  using tile_desc = typename payload_t::tile_desc;
+  using load_dtype = typename payload_t::mem_dtype;
+  constexpr uint32_t load_elems = payload_t::simd_exec_size;
+  constexpr uint32_t pack_factor = payload_t::pack_factor;
+
+#pragma unroll
+  for (uint32_t i = 0; i < tile_desc::num_block_y; i++) {
+    uint32_t offset_y = i * tile_desc::block_size_y;
+#pragma unroll
+    for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
+      uint32_t offset_x = j * tile_desc::block_size_x;
+      auto reg_sub = tile.reg.xetla_select<tile_desc::block_elems, 1>(
+          (i * tile_desc::num_block_x + j) * tile_desc::block_elems);
+#pragma unroll
+      for (uint32_t sub_block_y = 0; sub_block_y < tile_desc::block_size_y;
+           sub_block_y += 1) {
+        xetla_vector<load_dtype, load_elems> reg_tmp = 0;
+        uint32_t address_offset = payload_t::trans
+            ? offset_x * payload.pitch_in_bytes +
+                (offset_y + sub_block_y) * sizeof(dtype)
+            : offset_x * sizeof(dtype) +
+                (offset_y + sub_block_y) * payload.pitch_in_bytes;
+        reg_tmp = xetla_load_global<
+            load_dtype,
+            payload_t::simd_exec_size,
+            data_size::default_size,
+            L1,
+            L3>(payload.base_ptr, payload.base_offset + address_offset);
+
+        reg_sub
+            .xetla_select<load_elems * pack_factor, 1>(
+                sub_block_y * tile_desc::block_size_x)
+            .xetla_format<load_dtype>() = reg_tmp;
+      }
+    }
+  }
+
+  if constexpr (payload_t::trans) {
+    SW_BARRIER();
+    tile_transpose(tile);
+  }
+  if constexpr (payload_t::mem_transform) {
+    SW_BARRIER();
+    vnni_convert(tile);
+  }
+}
 /// @brief This function loads data from unaligned-2D memory surface.
 /// Loads an array of rectangular regions (X,Y)..(X+W,Y+H) from memory into
 /// registers. Each block will be loaded serially by its corresponding payload.

include/subgroup/tile/impl/op_function.hpp

@@ -676,4 +676,35 @@ layout_convert(T_dst& dst, T_src& src) {
     }
   }
 }
+
+template <typename T>
+void dump_mat(
+    T mat,
+    size_t tile_x = T::tile_size_x,
+    size_t tile_y = T::tile_size_y) {
+#pragma unroll
+  for (size_t row = 0; row < tile_y; row++) {
+#pragma unroll
+    for (size_t col = 0; col < tile_x; col++) {
+      sycl::ext::oneapi::experimental::printf(
+          "%d ", (int)(sycl::half)mat.reg[row * tile_x + col]);
+    }
+    sycl::ext::oneapi::experimental::printf("\n ");
+  }
+  sycl::ext::oneapi::experimental::printf("\n ");
+}
+template <typename T>
+void dump_mat_reg(T mat, size_t tile_x, size_t tile_y) {
+#pragma unroll
+  for (size_t row = 0; row < tile_y; row++) {
+#pragma unroll
+    for (size_t col = 0; col < tile_x; col++) {
+      sycl::ext::oneapi::experimental::printf(
+          "%d ", (int)(sycl::half)mat[row * tile_x + col]);
+    }
+    sycl::ext::oneapi::experimental::printf("\n ");
+  }
+  sycl::ext::oneapi::experimental::printf("\n ");
+}
+
 } // namespace gpu::xetla::subgroup

include/subgroup/tile/impl/payload_xe.hpp

@@ -404,14 +404,15 @@ struct mem_payload_t<
   // for pvc, we can use simd16 or simd32
   static constexpr uint32_t min_store_bytes = 16 * sizeof(dtype);
   static constexpr uint32_t max_store_bytes = 32 * sizeof(dtype);
-  static constexpr uint32_t num_channel =
+  static constexpr uint32_t simd_channel =
       ((tile_bytes % max_store_bytes) == 0 &&
        (block_bytes % max_store_bytes) == 0)
       ? 32
       : 16;
-
-  static constexpr uint32_t num_channel_x = block_size_x;
-  static constexpr uint32_t num_channel_y = num_channel / num_channel_x;
+  static constexpr uint32_t num_channel =
+      (simd_channel >= block_size_x) ? block_size_x : simd_channel;
+  static constexpr uint32_t num_channel_x = block_size_x; // 16
+  static constexpr uint32_t num_channel_y = num_channel / num_channel_x; // 1
   static constexpr uint32_t store_elems = num_channel_y * block_size_x;
 
   xetla_vector<uint32_t, num_channel> channel_offset;

tests/integration/data_transformer/common.hpp

@@ -122,6 +122,7 @@ class TestBase {
   using data_type_in = float;
   using data_type_out = bf16;
   using data_type_acc = float;
+  static constexpr gpu_arch gpu_arch = gpu_arch::XeHpc;
 };
 
 class Test_fp32tobf16_128_64 : public TestBase {