commit 4a034095698d5d0f3e49003fe9557c1aea7b6638
author Chip Kerchner <chip.kerchner@ibm.com> Mon Feb 13 21:30:45 2023 +0000
committer Rasmus Munk Larsen <rmlarsen@google.com> Mon Feb 13 21:30:45 2023 +0000
tree c36f81237d4d31d4ec1759b6994f5fa88382bb3e
parent 77b48c440ed1de189b5296feb86e24536f645065

Fix problem with array conversions BF16->F32 in Power.

Eigen/src/Core/arch/AltiVec/MatrixProductMMAbfloat16.h

diff --git a/Eigen/src/Core/arch/AltiVec/MatrixProductMMAbfloat16.h b/Eigen/src/Core/arch/AltiVec/MatrixProductMMAbfloat16.h
index 392e027..bdf4347 100644
--- a/Eigen/src/Core/arch/AltiVec/MatrixProductMMAbfloat16.h
+++ b/Eigen/src/Core/arch/AltiVec/MatrixProductMMAbfloat16.h
@@ -189,11 +189,43 @@
   ei_declare_aligned_stack_constructed_variable(float, result, cols*rows, 0);
 
   typedef typename DataMapper::LinearMapper LinearMapper;
+  Packet8us z = pset1<Packet8us>(0);
   for(Index j = 0; j < cols; j++){
     const LinearMapper res2 = res.getLinearMapper(0, j);
     float *result2 = result + j*rows;
-    BFLOAT16_UNROLL
-    for(Index i = 0; i < rows; i++){
+    Index i = 0;
+    for(; i + 32 <= rows; i+=32){
+      Packet8us r32_0 = res2.template loadPacket<Packet8bf>(i +  0).m_val;
+      Packet8us r32_1 = res2.template loadPacket<Packet8bf>(i +  8).m_val;
+      Packet8us r32_2 = res2.template loadPacket<Packet8bf>(i + 16).m_val;
+      Packet8us r32_3 = res2.template loadPacket<Packet8bf>(i + 24).m_val;
+      pstore(result2 + i +  0, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_0)));
+      pstore(result2 + i +  4, reinterpret_cast<Packet4f>(vec_mergel(z, r32_0)));
+      pstore(result2 + i +  8, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_1)));
+      pstore(result2 + i + 12, reinterpret_cast<Packet4f>(vec_mergel(z, r32_1)));
+      pstore(result2 + i + 16, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_2)));
+      pstore(result2 + i + 20, reinterpret_cast<Packet4f>(vec_mergel(z, r32_2)));
+      pstore(result2 + i + 24, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_3)));
+      pstore(result2 + i + 28, reinterpret_cast<Packet4f>(vec_mergel(z, r32_3)));
+    }
+    for(; i + 16 <= rows; i+=16){
+      Packet8us r32_0 = res2.template loadPacket<Packet8bf>(i +  0).m_val;
+      Packet8us r32_1 = res2.template loadPacket<Packet8bf>(i +  8).m_val;
+      pstore(result2 + i +  0, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_0)));
+      pstore(result2 + i +  4, reinterpret_cast<Packet4f>(vec_mergel(z, r32_0)));
+      pstore(result2 + i +  8, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_1)));
+      pstore(result2 + i + 12, reinterpret_cast<Packet4f>(vec_mergel(z, r32_1)));
+    }
+    for(; i + 8 <= rows; i+=8){
+      Packet8us r32_0 = res2.template loadPacket<Packet8bf>(i +  0).m_val;
+      pstore(result2 + i +  0, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_0)));
+      pstore(result2 + i +  4, reinterpret_cast<Packet4f>(vec_mergel(z, r32_0)));
+    }
+    for(; i + 4 <= rows; i+=4){
+      Packet8us r32_0 = res2.template loadPacketPartial<Packet8bf>(i +  0, 4).m_val;
+      pstore(result2 + i +  0, reinterpret_cast<Packet4f>(vec_mergeh(z, r32_0)));
+    }
+    for(; i < rows; i++){
       result2[i] = res2(i);
     }
   }