ทำไมการคูณ Matrix ให้เร็วปฏิวัติ Deep Learning อย่างไร

เราเคยคุยเรื่องนี้กับเพื่อนสมัยที่ยังเรียน ป.ตรี อยู่ว่า แค่คณิตศาสตร์ ม.4 อย่างการคูณ Matrix ให้ได้ มันทำให้งานด้าน Deep Learning มันก้าวหน้าไปไกลมาก ๆ วันนี้เราจะมาเล่าให้อ่านกันว่า ทำไมมันถึงเป็นแบบนั้น และตอนนี้เราทำได้เร็วที่สุดเท่าไหร่

ทำไม Deep Learning ใช้ Matrix

จริง ๆ เราจะแค่ Deep Learning ไม่ได้หรอก จริง ๆ มันใช้ตั้งแต่แค่เราทำ Neural Network เล็ก ๆ ก็ใช้ ภาพที่เราเข้าใจกับการทำ Neural Network คือ การมี Node ก้อนกลม ๆ ต่อ ๆ กันไปเป็นชั้น ๆ ไปเรื่อย ๆ เบสิกเลย ก็คือ มี 3 Layer คือ Input, Hidden และ Output ในส่วนของ Deep Learning คำว่า Deep ก็คือ การยัด Hidden Layer เยอะ ๆ เอาให้ลึก ๆ Deep สุด ๆ กันไปเลย (จริง ๆ มันมีเหตุผลในเชิงคณิตศาสตร์อยู่ แต่ไว้ค่อยว่ากัน)

float calNode (float *input, float *weight, float activation_fn (float)) {
  float sum = 0;
  for (int i=0; i<sizeof(input)/sizeof(float);i++) {
    sum += input[i] * weight[i];
  }
  return activation_fn(sum+bias)
}

ถ้าเราลองมองให้ลึกขึ้น ในแต่ละ Node การทำงานมันง่ายมาก ๆ คือ เราเอา Weight คูณด้วย Output ตัวนั้น ๆ และค่อยเอาผลลัพธ์ตรงนั้นไปใส่ Activation Function เข้าไป ซึ่งจะเป็นอะไรก็ว่ากันไป Code มันจะเป็นแบบด้านบน สั้น ๆ ง่าย ๆ เลย ถ้าเราลองดูดี ๆ เราจะเห็นว่า จริง ๆ แล้วสิ่งที่เราทำ กับ Input และ Weight มันคือ การหา Dot Product นั่นเอง

ที่เรายกตัวอย่างมา นี่มันคือแค่ของ Node เดียวเท่านั้น ลองคิดว่าใน Model สักตัว มันต้องมี Node กี่อัน มี Weight ที่ต่อเข้าไปเท่าไหร่ หากเราจะ Represent มันแยกเป็นก้อน ๆ แบบนั้นน่าจะยุ่งยากแน่นอน ทำให้ในความเป็นจริง เราจะ Represent มันในแต่ละ Layer ออกเป็น Matrix ขนาดใหญ่ ๆ ไป นั่นทำให้ เวลาเราทำงานจริง ๆ มันคือการทำ Matrix Operation กับ Matrix นั่นเอง ยิ่ง Model ของเรามีความซับซ้อนมากเท่าไหร่ ยิ่งทำให้เราต้องทำงานกับ Matrix ขนาดใหญ่มากขึ้น ส่งผลในเรื่องของเวลามากขึ้นเท่านั้น

Idea ของการทำ Representation ลักษณะนี้ ใช้กับ Neural Network Model ได้ทุกตัวเลยเด้อ ตั้งแต่ Perceptron ง่าย ๆ จนไปถึง Model อย่าง Recurrent Neural Network (RNN) และใหม่ ๆ อย่าง Attention Network ได้เลยละ เพราะมันมาจากพื้นฐานเดียวกันทั้งหมด แต่เปลี่ยนรูป ลักษณะภายใน และการเชื่อมต่อเท่านั้นเอง

CPU vs Neural Network

float calNode (float *input, float *weight, float activation_fn (float)) {
  float sum = 0;
  for (int i=0; i<sizeof(input)/sizeof(float);i++) {
    sum += input[i] * weight[i];
  }
  return activation_fn(sum+bias)
}

หากเราลองพิจารณา Code สำหรับการทำ Dot Product อีกครั้ง เราจะเห็นว่า มันใช้ O(n) เท่านั้นเอง นั่นแปลว่า เมื่อ Matrix เรามีขนาดใหญ่ขึ้นเรื่อย ๆ เวลาที่ใช้ในกรณีที่แย่ที่สุดจะเติบโตขึ้นเป็นเส้นตรงนั่นเอง

void matMulCPU (float *A, float *B, int N) {
  for (int i=0;i<N;i++) {
    for (int j=0;j,N;j++) {
      float currentValue = 0;
      for (int k=0;k<N;k++) {
        currentValue += A[i*N+k] * B[k*N+j]
      }
      A[i][j] = currentValue
    }
  }
}

อันที่เราคิดว่าหนักสุด ๆ คือ การคูณ Matrix หากเราคูณกันแบบ ม.ปลาย คือ หลัก คูณ แถว ไปเรื่อย ๆ พูดง่าย ๆ มันคือ Dot Product ของแถวและหลัก นี่คือสิ่งที่เกิดขึ้น เราจะเห็นว่า มันใช้ O(n^3) บอกเลยว่า ชิบหายการคอมพิวเตอร์สุด ๆ แน่นอน เพราะ CPU โดยทั่วไป มันทำงานเป็นแบบ Single Thread เท่ากับว่า การคำนวณพวกนี้จะเกิดขึ้นทีละตัว ๆ ไปเรื่อย ๆ

ถามว่า แล้วเราจะเร่งความเร็วในการทำงานได้อย่างใด โดยที่ไม่ต้องเพิ่มความเร็วสัญญาณนาฬิกา (Clock Speed) และ Instruction Per Clock (IPC) บน CPU

วิธีที่คิดเร็ว ๆ คือ การใช้ Multithread Programming เข้ามาช่วยได้นิ เพราะ แต่ละตำแหน่งของคำตอบไม่ได้เกี่ยวข้องกับตำแหน่งอื่น ๆ ประกอบกับเดี๋ยวนี้ CPU ตัวนึงเรามีสัก 8 Core หรือมากกว่าแล้ว เราก็กระจายงานไปตาม Core เช่น ตำแหน่งที่ 0 เราให้ Core 0, ตำแหน่งที่ 1 ให้ Core 1 ทำแบบนี้ไปเรื่อย ๆ เท่ากับว่า ถ้าเรามี 8 Core เราจะล่นเวลาที่ใช้ไปได้ สูงสุด 8 เท่าตัวเลยนะ ใช่ครับ มันทำแบบนั้นได้จริง ๆ แต่ยังไง มันก็ถือว่ายังน้อยมาก ๆ เมื่อเทียบกับขนาดของ Matrix หลักแสน

GPU vs Neural Network

เพื่อเป็นการแก้ปัญหานี้ เราเริ่มมีการเอา GPU เข้ามาใช้ในงานพวก Neural Network กัน ความต่างของ CPU, GPU และ NPU เราเคยเล่าไปแล้วในวีดีโอด้านบนนี้

แต่ว่าสั้น ๆ คือ ภายใน GPU มันเป็นหน่วยประมวลผลพิเศษที่ด้านในประกอบด้วย Core จำนวนมาก เรากำลังพูดถึงตัวเลขหลักหลายพันจนถึงหมื่นเลยละ นั่นทำให้ เราสามารถคำนวณ ผลลัพธ์ของแต่ละตำแหน่งได้พร้อม ๆ กันหลักหมื่นตำแหน่งได้เลย

ทำให้เกิดคำถามว่า แล้วทำไมเราถึงไม่เอา GPU Design มาใช้แทน CPU เลยละ เหตุผลมันอยู่ที่ว่า เจ้า Core ของ GPU ที่มันมีจำนวนมาก จริง ๆ มันเป็น Core ที่ไม่ได้มีอะไรมาก นอกจาก FPU และ INTU รวมกับพวก Dispatch Unit เท่านั้นเอง เรียกว่า 1 Core ของ GPU น่าจะเทียบได้กับ ALU ที่อยู่ภายใน Core ของ CPU เท่านั้นเอง ขาดองค์ประกอบอื่น ๆ ที่ใน 1 CPU Core พึงมี เช่น Control Unit และ Memory Unit โดยทั้งสองอย่างที่ขาดหายไป มันจะถูกแชร์กันใน กลุ่มก้อนของ GPU Core เช่น เรามีกลุ่มของ Core อยู่ในกลุ่มนั้นอาจจะมี 32 Core เลยก็ได้ เขาออกแบบมาในลักษณะของ SIMD นั่นแปลว่า GPU Core ที่อยู่ในกลุ่มเดียวกันจะสามารถรันคำสั่งเดียวกันได้เท่านั้น แลกมากับการทำงานแบบ Parallel ที่ให้ Throughput สูงมาก ๆ

แต่หากมี Branching เกิดขึ้น ในสัก Core นึง จะต้องมีคนใดคนหนึ่งหยุดรอ ซึ่งจะแตกต่างจาก CPU Core ที่แต่ละ Core เขามีองค์ประกอบที่ทำให้สามารถทำงานแยกกันได้ เหมือนมีหลายคนทำงานที่แตกต่างกันพร้อมกันได้ ทำให้ CPU จะได้เปรียบเรื่องนี้มากกว่า ซึ่งงานส่วนใหญ่ มันไม่ได้เหมือนกันอยู่แล้ว และมันมี Branching สารพัด ทำให้ เราไม่สามารถเอาลักษณะของการออกแบบ GPU มาใช้แทน CPU ได้นั่นเอง

คิดภาพง่าย ๆ ว่า GPU Core มันเหมือนคนพายเรือ ที่ใช้คนจำนวนมากพายไปพร้อม ๆ กัน ถ้าต้องเปลี่ยนไปดูทาง ทุกคนก็ต้องดูทาง ไม่มีคนพาย แต่ CPU มีคนน้อยกว่า กลับสามารถพายเรือ พร้อมกับดูทางได้พร้อม ๆ กัน

How to do that?? -> CUDA Programming

ทีนี้ เราจะเขียนโปรแกรมเพื่อเข้าถึงการทำงานของ GPU ได้อย่างไรละ โดยปกติ GPU เขาจะมี API ในการเข้าถึงมาให้เราอยู่แล้ว วันนี้เราขอใช้ฝั่งที่เราทำงานคือ Nvidia GPU และภาษา C เขามี CUDA C API มาให้เราใช้งานได้

เรื่องตลกของ Post นี้คือ Code ทั้งหมดที่เขียนนี้ ก๊อปมาจากอันที่เขียนสมัยเรียน ป.ตรีทั้งนั้นเลย ทำให้สงสัยนะว่า เมื่อก่อนกรูเรียนอะไรพวกนี้รู้เรื่องได้ไงวะ โตมาพอต้องใช้ เออ น่าจะตั้งใจเรียนวิชา Parallel เนอะ

#include <stdio.h>
#include <stdlib.h>

int main () {
  // Init square matrix size
  int matrixSize = 10000;

  //Allocate memory for matrix A,B and R
  float* A = (float*) malloc(matrixSize*matrixSize*sizeof(float));
  float* B = (float*) malloc(matrixSize*matrixSize*sizeof(float));
  float* R = (float*) malloc(matrixSize*matrixSize*sizeof(float));
  
  // Random Sample Matrix of A and B
  for (int i=0;i<N;i++) {
    for (int j=0;i<N;j++) {
      A[i*matrixSize*j] = (float)(rand() / rand());
      B[i*matrixSize*j] = (float)(rand() / rand());
    }
  }

  // Alocate VRAM
  
  // Copy A,B into VRAM
  
  // Compute Multiplication in GPU
  
  // Copy C (Result Matrix) back to RAM

  // Free GPU Memory
  
  return 0;
}

การทำงานคร่าว ๆ เราแบ่งออกเป็น 6 ขั้นตอนใหญ่ ๆ ด้วยกัน คือ เราโหลด Matrix เข้าไปใน RAM ผ่าน CPU จากนั้น เราจะเริ่มไปขอ Allocate พื้นที่บน VRAM ที่อยู่บน GPU จากนั้นเราจะเอา Matrix ที่อยู่ใน GPU ย้ายไปที่ VRAM แล้วจึงเริ่มคำนวณ ย้ายผลลัพธ์กลับมาที่ RAM อีกครั้งเพื่อให้ CPU สามารถเข้าถึงได้เหมือนเดิม และเมื่อเราจอง Memory แล้วเราก็ต้องคืนกลับด้วย

// Allocate VRAM

// GPU Array Pointers
float* d_A, d_B, d_C;

// VRAM Allocation
cudaMalloc((void**) &d_A, matrixSize*matrixSize*sizeof(float));
cudaMalloc((void**) &d_B, matrixSize*matrixSize*sizeof(float));
cudaMalloc((void**) &d_C, matrixSize*matrixSize*sizeof(float));

เริ่มจากการ ขอ Allocate VRAM บน GPU กันก่อน โดยคำสั่งที่ใช้ คล้ายกับเวลาเราขอ Allocate Memory ปกติ ซึ่งเขาจะใช้คำสั่งที่ชื่อว่า cudaMalloc โดยมันจะรับ 2 Parameter เข้าไปคือ Pointer สำหรับชี้ไปที่ GPU Memory และขนาดที่เราต้องการจะจอง

แต่อาจจะสังเกตว่า ทำไมใน Parameter เราไม่ใส่ d_A เข้าไปตรง ๆ ทำไมเราจะต้องยัดเป็น Address ของ Pointer เอาละ เพราะ Function พวกนี้มันมี Return Type เป็น Void เลย สิ่งที่มันทำคือ มันจะคืน Address บน GPU กลับมาโดยการใส่เข้าไปเป็น Value ใน Address ของตัวแปร Pointer ที่เราใส่เข้าไปนั่นเอง

แต่ใน Code ด้านบนนี้เราไม่ได้มีการเช็ค Error นะว่า หากมันไม่สามารถ Allocate ได้ มันจะผ่านไปเลย ซึ่งโดยปกติ เวลาเราเขียนโปรแกรม เราควรจะเช็คเรื่องนี้สักหน่อยเด้อ ลองไปหาวิธีการใน Document ของ CUDA ได้

// Copy A,B into VRAM
cudaMemcpy(d_A, A, matrixSize*matrixSize*sizeof(float), cudaMemcpyHostToDevice);
cudaMemcpy(d_B, B, matrixSize*matrixSize*sizeof(float), cudaMemcpyHostToDevice);

ขั้นตอนต่อไปคือ เราจะทำการ Copy Matrix A และ B ที่เราได้ทำการ Generate ไว้เข้าไปที่ VRAM ใน CUDA API เขามีคำสั่งชื่อว่า cudaMemcpy สำหรับการย้ายข้อมูลอยู่ โดยเราจะต้องป้อนทั้งหมด 4 Parameter คือ Pointer ของปลายทาง, Pointer ต้นทาง, ขนาดที่ต้องการ Copy และ ทิศทางในการย้าย

โดยจะเห็นจาก Code ด้านบนว่า เรากำหนดปลายทางเป็น Pointer ที่ชี้ไปที่ตำแหน่งใน GPU จากต้นทางที่เป็น Matrix ที่เรา Generate ขึ้นมา ส่วนขนาดก็ตามปกติ คือ เอาแถวคูณหลักคูณด้วยขนาดของ Float สุดท้าย ทิศทาง เราใส่เป็น cudaMemcpyHostToDevice มันคือการคัดลอกจาก CPU ไปที่ GPU

// Compute Multiplication in GPU
dim3 dim_block(2,2,1);
dim3 dim_grid(ceil(N/2.0), ceil(N/2.0), 1);

cudaMatMul_Kernel <<< dim_grid, dim_block>>>(d_A, d_B, d_C,N);

มาถึงส่วนที่โคตรสนุกละ ภายใน GPU เขาจะมีหลาย Multiprocessing Unit เช่น ใน GPU รุ่น A อาจจะมี 4 Multiprocessing Unit และในแต่ละ Unit เราสามารถสร้าง Thread ได้ 100 ตัว นั่นแปลว่า จำนวน Thread ที่เราสามารถรันได้พร้อมกันคือ 4 x 100 = 400 Threads ในเวลาเดียวกัน ในการจัดสรร CUDA เขากำหนดมาให้ เราจัดงานทั้งหมดออกเป็น Block อาจจะเป็น 1D, 2D และ 3D ก็ได้ ในที่นี้เราขอให้ 2D แบบปลอม ๆ ละกัน โดยที่ แต่ละ Block มันจะถูกจัดเข้ากับ Multiprocessing Unit สักตัวนึง หากเรามีจำนวน Block ที่ต้องทำงานมากกว่าจำนวน Multiprocessing Unit มันก็จะเข้าคิวกันทำงานโดยอัตโนมัติ

อ่านแล้วอาจจะเข้าใจยาก คิดง่าย ๆ ว่า เราจะแยกงานออกมาเป็นหลาย ๆ ก้อน ในแต่ละก้อนนั้นจะถูกจัดสรรด้วยคนงานจำนวนหนึ่ง เช่นในที่นี้ จาก Code ด้านบน เราจะขอแบ่งให้ในแต่ละ Block มีขนาด 2 x 2 x 1 ทำให้ใน 1 Block จะต้อง เปิด Thread ทั้งหมด 2x2x1 = 4 Threads ด้วยกัน และสุดท้าย เราค่อย ๆ โยนแต่ละ Block เข้าไปพร้อมกับคำสั่งที่มันต้องทำ เราเรียก Function อันนี้ว่า Kernel หากใครรู้จักคำว่า SIMD ใช่ฮะ ตอนนี้เรากำลังทำแบบนั้นเลย

__global__ void cudaMatMul_kernel (float* d_A, float* d_B, float* d_C, int N) {
  // identify where am i
  int i = blockDim.y * blockDim.y + threadIdx.y;
  int j = blockDim.x * blockDim.x + threadIdx.x;

  // Check border
  if (i < N && j < N) {
    float currentValue = 0.0;
    for (int k=0; k<N;k++) {
      currentValue += d_A[i*N+k] * d_B[k*N+j];
    }

    d_C[i*N+j] = currentValue;
  }
}

ภายใน Kernel เราเริ่มจากการหาก่อนว่า ตัวเองอยู่ตรงไหน โดยการเช็คจาก เลข Block และ Thread เราก็จะพอรู้ละว่า เราจะต้องทำงานกับตำแหน่งไหน และเพื่อความปลอดภัย เราจะเช็คว่า เราอยู่ที่ขอบของ Matrix ที่เกินมาหรือไม่ เผื่อเวลามันปัดจำนวน Block ขึ้นตามที่เราสั่ง มันจะมีขอบที่มันไม่ต้องคำนวณ เราก็ข้ามไปได้เลย และสุดท้าย เราก็แค่หา Dot Product ปกติแล้วก็เอาค่าไปใส่ใน d_C ที่เก็บผลลัพธ์ของเราก็เป็นอันเรียบร้อย

// Copy C (Result Matrix) back to RAM
cudaMemcpy(C, d_C, matrixSize*matrixSize*sizeof(float), cudaMemcpyDeviceToHost);

หลังจากเราคำนวณเรียบร้อยแล้ว ผลลัพธ์ทั้งหมดอยู่ใน VRAM ของ GPU ที่เราสามารถเข้าถึงได้ในตัวแปร d_C แต่เราต้องการย้ายมันกลับมาที่ CPU เพื่อเอาไปใช้งานต่อ ดังนั้น เราจะต้องคัดลอกมันกลับมา

เราทำเหมือนเดิมคือ การใช้คำสั่ง cudaMemcpy แต่รอบนี้เราจะต้องกลับด้านนะ รอบนี้ปลายทางคือ C ที่อยู่ใน RAM จากต้นทาง d_C ที่อยู่ใน VRAM ขนาดเท่าเดิม และทิศทางต้องเปลี่ยนเป็น DeviceToHost หรือ จาก GPU ไปที่ CPU แทน

// Free GPU Memory
cudaFree(d_A);
cudaFree(d_B);
cudaFree(d_C);

สุดท้าย ท้ายสุด เมื่อเราจอง Memory และใช้งานเสร็จแล้ว ก็จงอย่าลืมคืน Memory กลับไปด้วย ไม่งั้น Memory Leak แตกกันทั้งวงแน่นอน เรียกได้ว่า น้ำแตกแยกทางเลยทีเดียว

สรุป

กลายเป็นว่า บทความนี้เหมือนมาสอน CUDA Programming เลอ แต่ประเด็นที่เราต้องการจะบอกคือ แค่การทำ Matrix Operation ที่เราเรียนกันตอนมัธยมมันมีประโยชน์กว่าที่เราคิด เมื่อเราเอามารวมกับความรู้ทางคอมพิวเตอร์ มันกลายเป็นรากฐานสำคัญของการพัฒนา AI สมัยใหม่เลยทีเดียว