From c3dd7f0a2ae5bbc2271a4bb9543976642ecfb46f Mon Sep 17 00:00:00 2001
From: "Fine, Gregory" <gregory.fine@intel.com>
Date: Wed, 2 Oct 2024 11:49:49 -0700
Subject: [PATCH] Fix load_2d inconsistency when reading non native types with
 VNNI transforms

---
 sycl/include/sycl/ext/intel/esimd/memory.hpp    | 11 +++++------
 sycl/test-e2e/ESIMD/lsc/lsc_load_2d_compare.cpp |  3 +++
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/sycl/include/sycl/ext/intel/esimd/memory.hpp b/sycl/include/sycl/ext/intel/esimd/memory.hpp
index bfdd13f50a337..9a8917504f011 100644
--- a/sycl/include/sycl/ext/intel/esimd/memory.hpp
+++ b/sycl/include/sycl/ext/intel/esimd/memory.hpp
@@ -4070,7 +4070,7 @@ __ESIMD_API simd<T, N> load_2d_impl(const T *Ptr, unsigned SurfaceWidth,
   uintptr_t Addr = reinterpret_cast<uintptr_t>(Ptr);
   constexpr lsc_data_order Transpose =
       Transposed ? lsc_data_order::transpose : lsc_data_order::nontranspose;
-  simd<RawT, ActualN> Raw =
+  simd<T, ActualN> Raw =
       __esimd_lsc_load2d_stateless<RawT, L1H, L2H, DS, Transpose, NBlocks,
                                    BlockWidth, BlockHeight, Transformed,
                                    ActualN>(Mask.data(), Addr, SurfaceWidth,
@@ -4096,17 +4096,16 @@ __ESIMD_API simd<T, N> load_2d_impl(const T *Ptr, unsigned SurfaceWidth,
     // +----+----+----+----+----+----+-----+-----+
     // * signifies the padded element.
 
-    simd<RawT, DstElements> Dst;
+    simd<T, DstElements> Dst;
 
     for (auto i = 0; i < NBlocks; i++) {
       auto DstBlock =
           Dst.template select<DstBlockElements, 1>(i * DstBlockElements);
 
       auto RawBlock = Raw.template select<GRFBlockSize, 1>(i * GRFBlockPitch);
-      DstBlock =
-          RawBlock.template bit_cast_view<RawT, GRFColSize, GRFRowPitch>()
-              .template select<GRFColSize, 1, GRFRowSize, 1>(0, 0)
-              .template bit_cast_view<RawT>();
+      DstBlock = RawBlock.template bit_cast_view<T, GRFColSize, GRFRowPitch>()
+                     .template select<GRFColSize, 1, GRFRowSize, 1>(0, 0)
+                     .template bit_cast_view<T>();
     }
 
     return Dst;
diff --git a/sycl/test-e2e/ESIMD/lsc/lsc_load_2d_compare.cpp b/sycl/test-e2e/ESIMD/lsc/lsc_load_2d_compare.cpp
index cc22cf918e349..d99719c71b712 100644
--- a/sycl/test-e2e/ESIMD/lsc/lsc_load_2d_compare.cpp
+++ b/sycl/test-e2e/ESIMD/lsc/lsc_load_2d_compare.cpp
@@ -79,12 +79,15 @@ int main() {
   result |= test<uint16_t>();
   result |= test<uint8_t>();
   result |= test<sycl::half>();
+  result |= test<bf16>();
 
   result |= test<float, true>();
   result |= test<uint32_t, true>();
 
   result |= test<uint16_t, false, true>();
   result |= test<uint8_t, false, true>();
+  result |= test<sycl::half, false, true>();
+  result |= test<bf16, false, true>();
 
   std::cout << (result ? "FAILED" : "passed") << std::endl;
   return 0;