Improved stored

Author: MarioSieg

python/benchmark/throughput_single.py

@@ -6,28 +6,28 @@
 import torch, piquant
 torch.set_num_threads(mp.cpu_count())
 
-TOTAL_GIB = 64
+TOTAL_GIB = 32
 SRC_DTYPE = torch.bfloat16
-Q_DTYPE   = torch.quint4x2
+Q_DTYPE = torch.quint4x2
 OUT_DTYPE = torch.bfloat16
 
-SCALE, ZP = 1.0, 2
-
-B_PER_ELEM  = torch.tensor([], dtype=SRC_DTYPE).element_size()  # 2 for bf16
+B_PER_ELEM = torch.tensor([], dtype=SRC_DTYPE).element_size()
 TOTAL_BYTES = TOTAL_GIB * (1 << 30)
-N_ELEMS     = TOTAL_BYTES // B_PER_ELEM
+N_ELEMS = TOTAL_BYTES // B_PER_ELEM
+print(f'Numel = {N_ELEMS}')
 
 with torch.inference_mode():
     x = torch.rand(N_ELEMS, dtype=SRC_DTYPE)
+    SCALE, ZP = piquant.torch.compute_quant_params(x, dtype=Q_DTYPE)
 
 t0 = time.perf_counter()
 with torch.inference_mode():
     q = piquant.torch.quantize(x, scale=SCALE, zero_point=int(ZP), dtype=Q_DTYPE)
 t1 = time.perf_counter()
-print(f"Quant time (bf16 -> uint4) for 100 GiB: {t1 - t0:.3f} s")
+print(f"Quant time (bf16 -> uint4) for {TOTAL_GIB} GiB: {t1 - t0:.3f} s")
 
 t0 = time.perf_counter()
 with torch.inference_mode():
     y = piquant.torch.dequantize(q, scale=SCALE, zero_point=int(ZP), dtype=OUT_DTYPE)
 t1 = time.perf_counter()
-print(f"Dequant time (uint4 -> bf16) for 100 GiB source: {t1 - t0:.3f} s")
+print(f"Dequant time (uint4 -> bf16) for {TOTAL_GIB} GiB source: {t1 - t0:.3f} s")

python/quant_benchmark.png

@@ -1189,14 +1189,25 @@ static auto PIQUANT_HOT dequant_uint4_to_bf16(
                     vf31 = _mm512_add_ps(vf31, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(o+i+112))), 16)));
                 #endif
             }
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+0), cvt_ps_to_bf16(vf00));
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+16), cvt_ps_to_bf16(vf10));
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+32), cvt_ps_to_bf16(vf20));
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+48), cvt_ps_to_bf16(vf30));
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+64), cvt_ps_to_bf16(vf01));
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+80), cvt_ps_to_bf16(vf11));
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+96), cvt_ps_to_bf16(vf21));
-            _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+112), cvt_ps_to_bf16(vf31));
+            #ifdef __AVX512BF16__
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+0),   std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf00)));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+16),  std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf10)));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+32),  std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf20)));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+48),  std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf30)));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+64),  std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf01)));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+80),  std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf11)));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+96),  std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf21)));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+112), std::bit_cast<__m256i>(_mm512_cvtneps_pbh(vf31)));
+            #else
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+0), cvt_ps_to_bf16(vf00));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+16), cvt_ps_to_bf16(vf10));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+32), cvt_ps_to_bf16(vf20));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+48), cvt_ps_to_bf16(vf30));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+64), cvt_ps_to_bf16(vf01));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+80), cvt_ps_to_bf16(vf11));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+96), cvt_ps_to_bf16(vf21));
+                _mm256_storeu_si256(reinterpret_cast<__m256i*>(o+i+112), cvt_ps_to_bf16(vf31));
+            #endif
         }
     #elif defined(__AVX2__)
     #elif defined(__SSE4_2__)