לדלג לתוכן

6.1 cuBLAS אלגברה לינארית מואצת פתרון

פתרון - cuBLAS - אלגברה לינארית מואצת

הערה: הזמנים, ה-GFLOP/s ואחוזי ה-peak בפלטים למטה הם דוגמאות ממכונת H100 SXM, וישתנו אצלכם לפי הכרטיס וגרסת ה-CUDA. מה שלא ישתנה הוא המבנה, סדרי הגודל, יחסי ה-speedup בין המסלולים, וסוגי השגיאות - עליהם התרגיל בנוי. כל הקטעים מניחים שהכותרות, מאקרו CUDA_CHECK ו-CUBLAS_CHECK, וכלי הייחוס/מדידה מ-הכנה נמצאים בראש הקובץ. הcompilation: nvcc -O2 -arch=sm_90a cublas_gemm.cu -o cublas_gemm -lcublas.

פתרון תרגיל 1 - כפל row-major נכון עם cublasSgemm

int main(void) {
    int N = 512;
    size_t bytes = (size_t)N * N * sizeof(float);
    float *hA = (float*)malloc(bytes), *hB = (float*)malloc(bytes);
    float *hC = (float*)malloc(bytes), *hRef = (float*)malloc(bytes);
    for (int i = 0; i < N*N; i++) {
        hA[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f;
        hB[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f;
    }

    float *dA, *dB, *dC;
    CUDA_CHECK(cudaMalloc(&dA, bytes));
    CUDA_CHECK(cudaMalloc(&dB, bytes));
    CUDA_CHECK(cudaMalloc(&dC, bytes));
    CUDA_CHECK(cudaMemcpy(dA, hA, bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dB, hB, bytes, cudaMemcpyHostToDevice));

    cublasHandle_t handle;
    CUBLAS_CHECK(cublasCreate(&handle));
    const float alpha = 1.0f, beta = 0.0f;

    // (4) the correct trick: request C^T = B^T * A^T -> dB before dA
    CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                             N, N, N, &alpha,
                             dB, N,          // first operand = our B
                             dA, N,          // second operand = our A
                             &beta, dC, N));
    CUDA_CHECK(cudaMemcpy(hC, dC, bytes, cudaMemcpyDeviceToHost));

    // (5) verification against the CPU reference
    matmulCPU(hA, hB, hRef, N);
    printf("[correct order dB,dA]\n");
    checkResult(hRef, hC, N, 1e-3f);

    // (6) the wrong order: dA before dB -> computes (B*A)^T, wrong
    CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                             N, N, N, &alpha,
                             dA, N, dB, N, &beta, dC, N));
    CUDA_CHECK(cudaMemcpy(hC, dC, bytes, cudaMemcpyDeviceToHost));
    printf("[wrong order dA,dB]\n");
    checkResult(hRef, hC, N, 1e-3f);

    cublasDestroy(handle);
    return 0;
}

פלט צפוי:

[correct order dB,dA]
  max relative error = 4.170e-04 (tolerance 1.0e-03)  -> PASS
[wrong order dA,dB]
  max relative error = 1.988e+00 (tolerance 1.0e-03)  -> FAIL

הסדר הנכון עובר בקלות (השגיאה כ-4e-4 היא רק הפרש הצבירה FP32 מול double בייחוס). הסדר השגוי מחזיר תוצאה רחוקה לגמרי - ובלי שום שגיאת cuBLAS, כי הממדים חוקיים. מה שקיבלנו הוא למעשה (B*A)^T, כפל אחר לגמרי.

למה זה עבד: buffer של מטריצה row-major זהה סיבית-לסיבית ל-transpose שלה שנקרא column-major. לכן כשמבקשים מ-cuBLAS B^T * A^T (בכך שמעבירים את dB לפני dA), התוצאה ב-buffer של C היא (A*B)^T ב-column-major, שהיא בדיוק A*B שלנו כשקוראים row-major. הסדר ההפוך מחשב (B*A)^T - כפל תקין מבחינת cuBLAS, שגוי מבחינתנו.

איך להכליל: בכל קריאת cuBLAS ממטריצות row-major, החליפו את סדר האופרנדים (הימני קודם) ובחרו m=מספר עמודות הפלט, n=מספר שורות הפלט. cuBLAS לעולם לא תזהיר על פריסה שגויה - היא רואה רק בתים ו-ld; האחריות לפריסה הנכונה כולה עליכם, ולכן בדיקת נכונות מול ייחוס היא חובה בקריאה הראשונה.

פתרון תרגיל 2 - cuBLAS מול הקרנל שלכם והפער מהתקרה

// timer identical to 3.5/5.3
template <typename LaunchFn>
float timeIt(LaunchFn launch, int reps = 20) {
    cudaEvent_t s, e;
    CUDA_CHECK(cudaEventCreate(&s)); CUDA_CHECK(cudaEventCreate(&e));
    launch();                                   // warmup
    CUDA_CHECK(cudaDeviceSynchronize());
    CUDA_CHECK(cudaEventRecord(s));
    for (int r = 0; r < reps; r++) launch();
    CUDA_CHECK(cudaEventRecord(e));
    CUDA_CHECK(cudaEventSynchronize(e));
    float ms = 0; CUDA_CHECK(cudaEventElapsedTime(&ms, s, e));
    CUDA_CHECK(cudaEventDestroy(s)); CUDA_CHECK(cudaEventDestroy(e));
    return ms / reps;
}

// ... inside main, N = 4096, dA/dB/dC allocated and initialized ...
cublasHandle_t handle; CUBLAS_CHECK(cublasCreate(&handle));
const float alpha = 1.0f, beta = 0.0f;

auto cublasLaunch = [&]{
    cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N,
                &alpha, dB, N, dA, N, &beta, dC, N);
};
float msCublas = timeIt(cublasLaunch);
report("cublasSgemm", N, msCublas, H100_FP32_PEAK_GFLOPS);

// mmTiled from 3.5
dim3 block(32, 32);
dim3 grid((N + 31) / 32, (N + 31) / 32);
auto tiledLaunch = [&]{ mmTiled<<<grid, block>>>(dA, dB, dC, N); };
float msTiled = timeIt(tiledLaunch);
report("tiled(3.5)", N, msTiled, H100_FP32_PEAK_GFLOPS);

printf("  speedup cuBLAS over tiled = %.1fx\n", msTiled / msCublas);

פלט צפוי:

  cublasSgemm    N=4096     2.400 ms     57266.7 GFLOP/s  85.60% peak
  tiled(3.5)     N=4096     8.000 ms     17180.0 GFLOP/s  25.68% peak
  speedup cuBLAS over tiled = 3.3x

cuBLAS תופס ~86% מתקרת ה-FP32 של H100, בעוד ה-kernel ה-tiled שלנו נשאר על ~26%. הפער הוא פי 3.3.

למה זה עבד: cuBLAS בוחרת אוטומטית kernel שכתבו מומחי NVIDIA ברמת SASS, עם register tiling, double-buffering והיפוך פנימי של הלולאות המכוונן בדיוק ל-SM של Hopper. ה-tiled שלנו עצר ב-shared-memory tiling בלבד, שבו כל thread מחשב איבר פלט אחד; שאר הסולם (כל thread מחשב tile פלט מה-registers) עוד לפנינו בפרויקט 9.3.

איך להכליל: cuBLAS היא תקרת הייחוס לכל kernel שתכתבו. מדדו תמיד רק את זמן ה-kernel (בין events, בלי ה-cudaMemcpy), עם warmup, ובטאו את התוצאה כאחוז מהתקרה של אותו דיוק. הפער בין ה-kernel שלכם ל-cuBLAS הוא מפת הדרכים של האופטימיזציות שנשארו לכם.

פתרון תרגיל 3 - הדלקת Tensor Cores עם cublasGemmEx

// (1) TF32: the same FP32 input in memory, only the compute type changes
auto tf32Launch = [&]{
    cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N,
                 &alpha,
                 dB, CUDA_R_32F, N,
                 dA, CUDA_R_32F, N,
                 &beta,
                 dC, CUDA_R_32F, N,
                 CUBLAS_COMPUTE_32F_FAST_TF32,
                 CUBLAS_GEMM_DEFAULT);
};
float msTf32 = timeIt(tf32Launch);
report("GemmEx TF32", N, msTf32, 494700.0);   // H100's TF32 ceiling
// verification with relaxed tolerance: 1e-2

// (3) FP16: the data itself in half
__half *dAh, *dBh, *dCh;
CUDA_CHECK(cudaMalloc(&dAh, (size_t)N*N*sizeof(__half)));
CUDA_CHECK(cudaMalloc(&dBh, (size_t)N*N*sizeof(__half)));
CUDA_CHECK(cudaMalloc(&dCh, (size_t)N*N*sizeof(__half)));
// float -> half conversion kernel
f2h<<<(N*N+255)/256, 256>>>(dA, dAh, N*N);
f2h<<<(N*N+255)/256, 256>>>(dB, dBh, N*N);
CUDA_CHECK(cudaGetLastError());

auto fp16Launch = [&]{
    cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N,
                 &alpha,
                 dBh, CUDA_R_16F, N,
                 dAh, CUDA_R_16F, N,
                 &beta,
                 dCh, CUDA_R_16F, N,
                 CUBLAS_COMPUTE_32F,           // FP32 accumulation over FP16 input
                 CUBLAS_GEMM_DEFAULT);
};
float msFp16 = timeIt(fp16Launch);
report("GemmEx FP16", N, msFp16, 989400.0);   // H100's FP16 ceiling

ה-kernel הזעיר להמרת דיוק:

__global__ void f2h(const float* in, __half* out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) out[i] = __float2half(in[i]);
}

פלט צפוי (עם השורה של FP32 מתרגיל 2 להשוואה):

  cublasSgemm    N=4096     2.400 ms     57266.7 GFLOP/s  85.60% peak   (vs 66.9 TFLOPS)
  GemmEx TF32    N=4096     0.430 ms    319581.4 GFLOP/s  64.60% peak   (vs 494.7 TFLOPS)
  GemmEx FP16    N=4096     0.229 ms    600000.0 GFLOP/s  60.64% peak   (vs 989.4 TFLOPS)

  speedup: TF32 = 5.6x , FP16 = 10.5x  (vs cublasSgemm FP32)

בדיקות הנכונות:

  TF32  max relative error = 3.8e-03 (tolerance 1.0e-02)  -> PASS
  FP16  max relative error = 9.1e-03 (tolerance 2.0e-02)  -> PASS

למה זה עבד: cublasGemmEx מפרידה את טיפוס הנתונים מטיפוס החישוב. ב-TF32 הנתונים נשארים FP32 בזיכרון אבל הכפל רץ על ה-Tensor Cores בדיוק חתוך (19 ביט), ולכן זה "כמעט חינם" - פי 5.6 בלי לגעת בנתונים. ב-FP16 הנתונים עצמם רזים (16 ביט), מה שמכפיל שוב את התפוקה, אך הצבירה נשארת FP32 כדי שהשגיאה לא תתפוצץ - בדיוק ההוראה HMMA מ-1.5. ככל שהדיוק נמוך יותר, כך התפוקה גבוהה יותר וכך גם השגיאה היחסית, ולכן הרפינו את הסובלנות בכל צעד.

איך להכליל: המעבר מ-CUDA Cores ל-Tensor Cores הוא הקפיצה הגדולה ביותר בביצועי GEMM, והוא נעשה כאן בשינוי compute type בלבד. בחרו את הדיוק הנמוך ביותר שהיישום שלכם סובל (TF32 כמעט תמיד בטוח לאימון; FP16/BF16 עם צבירת FP32 לרוב הרשתות; FP8 להסקה מכומתת). תמיד השוו מול התקרה של אותו דיוק - להשוות FP16 מול תקרת FP32 הוא חסר משמעות.

פתרון תרגיל 4 - לשבור את ה-leading dimension ולתקן

שבירה א - ld קטן מדי (שגיאת ערך). משנים את ה-ld של dB מ-N ל-N/2:

cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N,
            &alpha, dB, N/2, dA, N, &beta, dC, N);   // ld=N/2 wrong

פלט צפוי:

cuBLAS error at cublas_gemm.cu:120
  'cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N, &alpha, dB, N/2, dA, N, &beta, dC, N)'
  -> CUBLAS_STATUS_INVALID_VALUE

מכיוון שהאופרנד הראשון ב-OP_N הוא בגודל m x k והוא צריך ld >= max(1,m) = N, ערך N/2 הוא בלתי-חוקי מפורשות. cuBLAS בודקת את הארגומנטים לפני ההרצה ומחזירה CUBLAS_STATUS_INVALID_VALUE מיד, וה-CUBLAS_CHECK עוצר את התוכנית בשורה המדויקת. זו התקלה ה"נעימה" - היא צועקת.

שבירה ב - ld שגוי אך חוקי (תוצאה שגויה בשקט). עם מטריצות מלבניות M=512, K=256, N=128 והקריאה הנכונה:

// row-major: A(M x K), B(K x N), C(M x N)
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
            N, M, K, &alpha,   // m=N, n=M, k=K
            dB, N,             // B^T: ld=N
            dA, K,             // A^T: ld=K
            &beta, dC, N);     // C^T: ld=N

פלט צפוי (הנכון):

[rect correct]
  max relative error = 5.02e-04 (tolerance 1.0e-03)  -> PASS

עכשיו מקלקלים רק את ה-ld של dA מ-K ל-M:

cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
            N, M, K, &alpha,
            dB, N,
            dA, M,             // <-- wrong: needs K, not M
            &beta, dC, N);

פלט צפוי:

[rect wrong ld on A]
  max relative error = 1.734e+00 (tolerance 1.0e-03)  -> FAIL

כאן M=512 הוא >= max(1,k)=256, ולכן ה-ld=M חוקי מבחינת cuBLAS אך שגוי סמנטית: cuBLAS קוראת מהצעד הלא נכון בזיכרון ומחזירה מספרים שגויים בשקט, בלי שום קוד שגיאה - זו בדיוק התקלה המסוכנת. לו היינו בוחרים M < K (למשל M=256, K=512), ה-ld=M היה קטן מ-max(1,k) והיינו מקבלים CUBLAS_STATUS_INVALID_VALUE במקום תוצאה שגויה בשקט. שני התרחישים חשובים: ld קטן מדי מ-max(1,m) נתפס כשגיאה; ld גדול מספיק אך לא תואם את הפריסה האמיתית עובר בשקט ומזהם את התוצאה.

התיקון. מחזירים כל ld לערך הנכון של התעלול: dB עם ld=N, dA עם ld=K, dC עם ld=N, ומקבלים שוב PASS.

למה זה עבד: ה-ld הוא הצעד בזיכרון בין שתי עמודות עוקבות ב-column-major, כלומר מספר השורות של המטריצה כפי ש-cuBLAS רואה אותה. בתעלול החלפת האופרנדים, dB הוא B^T בגודל N x K ולכן ld=N; dA הוא A^T בגודל K x M ולכן ld=K; dC הוא C^T בגודל N x M ולכן ld=N. ld קטן מ-m בלתי-אפשרי פיזית (עמודה לא נכנסת), ולכן נתפס; ld אחר-אך-גדול-מספיק פשוט מזיז את המצביע בצעדים שגויים.

איך להכליל: זכרו את הכלל - ב-OP_N ה-ld של אופרנד חייב להיות >= מספר השורות שלו כפי ש-cuBLAS רואה אותו. כשמשהו לא מסתדר, בדקו קודם את שלושת ה-ld: INVALID_VALUE פירושו ld קטן מדי (או ממד שלילי), ותוצאה שגויה בלי שגיאה פירושה כמעט תמיד ld שגוי-אך-חוקי או סדר אופרנדים הפוך. תמיד אמתו מול ייחוס אחרי כל שינוי בממדים או ב-ld.

פתרון תרגיל 5 (בונוס) - לראות את ה-HMMA ב-SASS ואת בחירת ה-kernel

nsys profile -o gemm --stats=true ./cublas_gemm

בטבלת cuda_gpu_kern_sum רואים שם kernel שונה לכל דיוק (דוגמה מקוצרת; השמות המדויקים משתנים בין גרסאות cuBLAS):

 Time(%)  Name
 -------  ---------------------------------------------------------------
   41.2   ampere_sgemm_128x128_nn                         (FP32 path)
   19.8   cutlass::Kernel<...tf32_...256x128...sm90...>    (TF32 path)
   12.4   cutlass::Kernel<...f16_...256x128...sm90...>     (FP16 path)

פרופיל של kernel יחיד עם Nsight Compute:

ncu --set full -k regex:".*gemm.*|.*Kernel.*" ./cublas_gemm

במדד ניצול ה-Tensor רואים את ההבדל המהותי:

  path         sm__pipe_tensor_op_hmma cycles active   conclusion
  ---------------------------------------------------------------
  FP32         ~0%                                      CUDA Cores only, no Tensor
  TF32         ~55%                                     Tensor Cores active (TF32)
  FP16         ~78%                                     Tensor Cores saturated (HMMA)

הרצת FP16 בשני גדלים שונים מראה ש-cuBLAS בוחרת kernels שונים:

  N=256    ->  ...f16_64x64_...     (small tile, fits a small problem)
  N=8192   ->  ...f16_256x128_...   (large tile, saturates Tensor Cores)

למה זה עבד: המדד sm__pipe_tensor_op_hmma מודד כמה מחזורים pipeline ה-Tensor Core היה פעיל - ולכן במסלול FP32 (שרץ על CUDA Cores) הוא כמעט אפס, ובמסלול FP16 הוא גבוה. אלו בדיוק הוראות ה-HMMA שפירקנו ב-1.5, רק שכאן cuBLAS ייצרה אותן עבורנו. הבחירה של tile שונה לגדלים שונים היא ההיוריסטיקה האטומה בפעולה: cuBLAS בחרה מתוך מאגר ה-kernels שלה את זה שמתאים לצורת הבעיה.

איך להכליל: nsys נותן את בחירת ה-kernel (איזה מימוש cuBLAS בחרה), ו-ncu נותן את התנהגותו (ניצול Tensor, memory-bound מול compute-bound, occupancy). כשתפרפלו GEMM ותראו שם kernel ארוך ומקודד עם hmma ב-SASS, תדעו בדיוק על מה אתם מסתכלים - וזה בדיוק המודל המנטלי שבנינו ב-1.5.