File size: 2,368 Bytes
e873e70 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | // Push CPU relu past 50 GB/s on i9-13900K: sweep thread count x per-thread ILP,
// AVX2 + non-temporal stores. Find the bandwidth-optimal config.
#include <cstdio>
#include <cstdint>
#include <cstdlib>
#include <immintrin.h>
#include <omp.h>
#include <chrono>
#include <algorithm>
// AVX2 NT-store relu, U vectors (8 floats each) per thread step.
template<int U>
static void relu_nt(float* __restrict out, const float* __restrict in, size_t n){
const __m256 z=_mm256_setzero_ps();
#pragma omp parallel for schedule(static)
for(size_t j=0;j<n;j+=(size_t)8*U){
for(int u=0;u<U;++u){
__m256 v=_mm256_loadu_ps(in+j+8*u);
_mm256_stream_ps(out+j+8*u, _mm256_max_ps(v,z));
}
}
_mm_sfence();
}
template<class F>
static double bench(F f, float* out, const float* in, size_t n){
for(int w=0;w<3;w++) f(out,in,n);
double best=1e30;
for(int r=0;r<8;r++){
auto t0=std::chrono::high_resolution_clock::now(); f(out,in,n);
auto t1=std::chrono::high_resolution_clock::now();
best=std::min(best,std::chrono::duration<double>(t1-t0).count());
}
return best;
}
int main(){
const size_t n=256ull*1024*1024; // 1 GB/array, divisible by 8*8
float* in =(float*)aligned_alloc(64,n*4);
float* out=(float*)aligned_alloc(64,n*4);
for(size_t i=0;i<n;++i) in[i]=((i&1)?-1.f:1.f)*float(i%97);
double gb=2.0*n*4/1e9;
// correctness of the unrolled NT kernel
for(size_t i=0;i<n;++i) out[i]=-7; relu_nt<4>(out,in,n);
bool ok=true; for(size_t i=0;i<n&&ok;i++){ float e=in[i]>0?in[i]:0; if(out[i]!=e) ok=false; }
printf("correctness: %s\n", ok?"PASS":"FAIL");
printf("AVX2 + NT-store, sweep threads x ILP (GB/s). DDR5 wall ~90 GB/s\n");
printf("threads | U=1 U=2 U=4 U=8\n");
int ts[]={4,8,12,16,24,32};
double bestbw=0; int bt=0,bu=0;
for(int t: ts){
omp_set_num_threads(t);
double b1=gb/bench(relu_nt<1>,out,in,n);
double b2=gb/bench(relu_nt<2>,out,in,n);
double b4=gb/bench(relu_nt<4>,out,in,n);
double b8=gb/bench(relu_nt<8>,out,in,n);
printf(" %2d | %5.0f %5.0f %5.0f %5.0f\n", t,b1,b2,b4,b8);
double m[4]={b1,b2,b4,b8}; int uu[4]={1,2,4,8};
for(int k=0;k<4;k++) if(m[k]>bestbw){bestbw=m[k];bt=t;bu=uu[k];}
}
printf(">> BEST: %.0f GB/s @ %d threads, U=%d (%.0f%% of ~90 GB/s wall)\n", bestbw,bt,bu,100*bestbw/90.0);
free(in); free(out);
return 0;
}
|