{"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n/**\n * @file vector_add.cu\n * \n * @brief We demonstrate adding 2 vectors of size 2048 using CUDA\n */\n\n/**\n * @brief CUDA kernel to perform vector addition\n * \n * @param d_v1 Device vector holding a input vector\n * @param d_v2 Device vector holding another input vector\n * @param d_v3 Device vector holding the sum of d_v1 and d_v2\n */\n__global__ void vector_add_kernel(float *d_v1, float *d_v2, float *d_v3, int n) {\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n if (i < n) {\n d_v3[i] = d_v1[i] + d_v2[i];\n }\n}\n\nint main() {\n int N = 2048;\n size_t size = sizeof(float) * N;\n\n // Allocate memory for the host vectors\n float *v1, *v2, *v3;\n cudaCheckError(::cudaMallocHost(&v1, size));\n cudaCheckError(::cudaMallocHost(&v2, size));\n cudaCheckError(::cudaMallocHost(&v3, size));\n\n // Fill the host vectors with data\n for (size_t i = 0; i < N; i++) {\n v1[i] = 6.4;\n v2[i] = 3.5;\n }\n\n // Create 3 device vectors for holding N float elements\n float *d_v1, *d_v2, *d_v3;\n cudaCheckError(::cudaMalloc(&d_v1, size));\n cudaCheckError(::cudaMalloc(&d_v2, size));\n cudaCheckError(::cudaMalloc(&d_v3, size));\n\n // Copy the host vectors into the device vectors\n cudaCheckError(::cudaMemcpy(d_v1, v1, size, cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(d_v2, v2, size, cudaMemcpyHostToDevice));\n\n // Launch the kernel using 8 thread blocks and 256 threads / block\n vector_add_kernel<<<8, 256>>>(d_v1, d_v2, d_v3, N);\n\n // Copy the result device vector into the result host vector\n cudaCheckError(::cudaMemcpy(v3, d_v3, size, cudaMemcpyDeviceToHost));\n\n cudaCheckError(::cudaDeviceSynchronize());\n\n // Verify the result host vector\n for (int i = 0; i < N; i++) {\n assert(v3[i] == (float)9.9);\n }\n\n // Free the device vector\n cudaCheckError(::cudaFree(d_v1));\n cudaCheckError(::cudaFree(d_v2));\n cudaCheckError(::cudaFree(d_v3));\n\n // Free the host vectors\n cudaCheckError(::cudaFreeHost(v1));\n cudaCheckError(::cudaFreeHost(v2));\n cudaCheckError(::cudaFreeHost(v3));\n\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/KernelAndLibExamples/vector_add/vector_add.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 2 // Rows of A, Rows of C\n#define N 2 // Columns of B, Columns of C\n#define K 3 // Columns of A, Rows of B\n\n// Kernel to perform matrix multiplication on a portion of the matrices\n__global__ void matMulKernelTP(int *A, int *B, int *C, int m, int n, int k, int col_start, int col_size) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int local_col = blockIdx.x * blockDim.x + threadIdx.x;\n int col = col_start + local_col;\n \n if (row < m && local_col < col_size) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * col_size + local_col] = value;\n }\n}\n\nint main() {\n // Host matrices\n int A[M][K] = { {1, 2, 3}, {4, 5, 6} };\n int B[K][N] = { {100, 200}, {300, 400}, {500, 600} };\n int C[M][N] = {0};\n\n // Calculate split for tensor parallelism across N dimension\n int cols_per_gpu = N / 2;\n \n // Arrays for each GPU\n int *d_A[2], *d_B[2], *d_C[2];\n \n // Set up each GPU\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n \n // Allocate memory on current GPU\n cudaCheckError(::cudaMalloc(&d_A[gpu], M * K * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_B[gpu], K * N * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_C[gpu], M * cols_per_gpu * sizeof(int)));\n \n // Copy input matrices to current GPU\n cudaCheckError(::cudaMemcpy(d_A[gpu], A, M * K * sizeof(int), cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(d_B[gpu], B, K * N * sizeof(int), cudaMemcpyHostToDevice));\n \n // Set grid and block dimensions for this GPU's portion\n dim3 threadsPerBlock(cols_per_gpu, M);\n dim3 numBlocks(\n (cols_per_gpu + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (M + threadsPerBlock.y - 1) / threadsPerBlock.y\n );\n \n // Launch kernel for this GPU's portion\n int col_start = gpu * cols_per_gpu;\n matMulKernelTP<<>>(\n d_A[gpu], d_B[gpu], d_C[gpu],\n M, N, K, col_start, cols_per_gpu\n );\n }\n \n // Copy results back from each GPU and combine\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n // Copy each row's portion separately to maintain correct layout\n for (int row = 0; row < M; row++) {\n cudaCheckError(::cudaMemcpy(\n &C[row][gpu * cols_per_gpu], // Destination in host matrix\n &d_C[gpu][row * cols_per_gpu], // Source from GPU\n cols_per_gpu * sizeof(int), // Size of this GPU's portion\n cudaMemcpyDeviceToHost\n ));\n }\n }\n \n // Output the result matrix\n std::cout << \"Result Matrix C:\" << std::endl;\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < N; j++) {\n std::cout << C[i][j] << \" \";\n }\n std::cout << std::endl;\n }\n \n // Free device memory on each GPU\n for (int gpu = 0; gpu < 2; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n cudaCheckError(::cudaFree(d_A[gpu]));\n cudaCheckError(::cudaFree(d_B[gpu]));\n cudaCheckError(::cudaFree(d_C[gpu]));\n }\n \n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/TensorParallelFromScratch/02_matmul_tp/matmul_tp_small.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 4096 // Rows of A, Rows of C\n#define N 8192 // Columns of B, Columns of C\n#define K 1024 // Columns of A, Rows of B\n#define NGPUS 2 // Number of GPUs to use for computation\n\n// Kernel to perform matrix multiplication on a portion of the matrices\n__global__ void matMulKernelTP(int *A, int *B, int *C, int m, int n, int k, int col_start, int col_size) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int local_col = blockIdx.x * blockDim.x + threadIdx.x;\n int col = col_start + local_col;\n \n if (row < m && local_col < col_size) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * col_size + local_col] = value;\n }\n}\n\n// Host function for matrix multiplication\nvoid matMulHost(const int* A, const int* B, int* C, int m, int n, int k) {\n for (int i = 0; i < m; i++) {\n for (int j = 0; j < n; j++) {\n int sum = 0;\n for (int p = 0; p < k; p++) {\n sum += A[i * k + p] * B[p * n + j];\n }\n C[i * n + j] = sum;\n }\n }\n}\n\nint main() {\n auto total_start = std::chrono::high_resolution_clock::now();\n // Host matrices\n int *A = new int[M * K];\n int *B = new int[K * N];\n int *C = new int[M * N](); // GPU result\n int *C_host = new int[M * N](); // CPU result\n\n // Initialize matrices with sample values\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < K; j++) {\n A[i * K + j] = (i + j) % 10; // Simple pattern to avoid large numbers in multiplication\n }\n }\n \n for (int i = 0; i < K; i++) {\n for (int j = 0; j < N; j++) {\n B[i * N + j] = (i + j) % 10;\n }\n }\n\n // Calculate split for tensor parallelism across N dimension\n const int cols_per_gpu = N / NGPUS;\n \n // Arrays for each GPU\n int *d_A[NGPUS], *d_B[NGPUS], *d_C[NGPUS];\n\n // Create streams and events per GPU\n cudaStream_t streams[NGPUS];\n cudaEvent_t startEvents[NGPUS], stopEvents[NGPUS];\n\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n cudaCheckError(::cudaStreamCreate(&streams[gpu]));\n cudaCheckError(::cudaEventCreate(&startEvents[gpu]));\n cudaCheckError(::cudaEventCreate(&stopEvents[gpu]));\n }\n\n // Set up each GPU\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n \n // Allocate memory on current GPU\n cudaCheckError(::cudaMalloc(&d_A[gpu], M * K * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_B[gpu], K * N * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_C[gpu], M * cols_per_gpu * sizeof(int)));\n \n // Copy input matrices to current GPU\n cudaCheckError(::cudaMemcpyAsync(d_A[gpu], A, M * K * sizeof(int), cudaMemcpyHostToDevice, streams[gpu]));\n cudaCheckError(::cudaMemcpyAsync(d_B[gpu], B, K * N * sizeof(int), cudaMemcpyHostToDevice, streams[gpu]));\n \n // Set grid and block dimensions for this GPU's portion\n const dim3 threadsPerBlock(16, 16); // Using 16x16 thread blocks for better occupancy\n const dim3 numBlocks(\n (cols_per_gpu + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (M + threadsPerBlock.y - 1) / threadsPerBlock.y\n );\n \n const int col_start = gpu * cols_per_gpu;\n\n cudaCheckError(::cudaEventRecord(startEvents[gpu], streams[gpu]));\n\n matMulKernelTP<<>>(\n d_A[gpu], d_B[gpu], d_C[gpu],\n M, N, K, col_start, cols_per_gpu\n );\n\n // Record stop event for this GPU\n cudaCheckError(::cudaEventRecord(stopEvents[gpu], streams[gpu]));\n }\n \n // Copy results back from each GPU and combine\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n // Copy each row's portion separately to maintain correct layout\n for (int row = 0; row < M; row++) {\n cudaCheckError(::cudaMemcpyAsync(\n &C[row * N + gpu * cols_per_gpu], // Destination in host matrix\n &d_C[gpu][row * cols_per_gpu], // Source from GPU\n cols_per_gpu * sizeof(int), // Size of this GPU's portion\n cudaMemcpyDeviceToHost,\n streams[gpu]\n ));\n }\n // Ensure all work in stream is completed before measuring timing.\n cudaCheckError(::cudaStreamSynchronize(streams[gpu]));\n }\n\n // Elapsed time per GPU.\n for (int gpu = 0; gpu < NGPUS; ++gpu) {\n float gpu_time = 0;\n cudaCheckError(::cudaEventElapsedTime(&gpu_time, startEvents[gpu], stopEvents[gpu]));\n std::cout << \"GPU \" << gpu << \" execution time: \" << gpu_time << \"ms\" << std::endl;\n }\n\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaEventDestroy(startEvents[gpu]));\n cudaCheckError(::cudaEventDestroy(stopEvents[gpu]));\n cudaCheckError(::cudaStreamDestroy(streams[gpu]));\n }\n\n // Start CPU timing\n auto cpu_start = std::chrono::high_resolution_clock::now();\n \n // Compute result on CPU for verification\n std::cout << \"Computing on CPU for verification...\" << std::endl;\n matMulHost(A, B, C_host, M, N, K);\n \n // Stop CPU timing\n auto cpu_end = std::chrono::high_resolution_clock::now();\n auto cpu_duration = std::chrono::duration_cast(cpu_end - cpu_start);\n \n // Print timing results\n std::cout << \"\\nTiming Results:\" << std::endl;\n std::cout << \"CPU Time: \" << cpu_duration.count() << \" ms\" << std::endl;\n \n // Compare results\n bool match = true;\n int mismatch_count = 0;\n for (int i = 0; i < M * N; i++) {\n if (C[i] != C_host[i]) {\n match = false;\n mismatch_count++;\n }\n }\n \n // Output comparison results\n std::cout << \"\\nResult comparison:\" << std::endl;\n if (match) {\n std::cout << \"GPU and CPU results match perfectly!\" << std::endl;\n } else {\n std::cout << \"Found \" << mismatch_count << \" mismatches between GPU and CPU results.\" << std::endl;\n }\n \n // Function to print a 4x4 region of both matrices\n auto print_region = [&](const char* region_name, int start_row, int start_col) {\n std::cout << \"\\n\" << region_name << \" (4x4):\" << std::endl;\n std::cout << \"GPU Result:\" << std::endl;\n for (int i = 0; i < 4; i++) {\n for (int j = 0; j < 4; j++) {\n std::cout << C[(start_row + i) * N + (start_col + j)] << \" \";\n }\n std::cout << std::endl;\n }\n \n std::cout << \"CPU Result:\" << std::endl;\n for (int i = 0; i < 4; i++) {\n for (int j = 0; j < 4; j++) {\n std::cout << C_host[(start_row + i) * N + (start_col + j)] << \" \";\n }\n std::cout << std::endl;\n }\n };\n \n // Print different regions of the matrices\n print_region(\"Top-left corner\", 0, 0);\n print_region(\"Middle region\", M/2 - 2, N/2 - 2);\n print_region(\"Bottom-right corner\", M - 4, N - 4);\n \n // Free device memory on each GPU\n for (int gpu = 0; gpu < NGPUS; gpu++) {\n cudaCheckError(::cudaSetDevice(gpu));\n cudaCheckError(::cudaFree(d_A[gpu]));\n cudaCheckError(::cudaFree(d_B[gpu]));\n cudaCheckError(::cudaFree(d_C[gpu]));\n }\n\n // Free host memory\n delete[] A;\n delete[] B;\n delete[] C;\n delete[] C_host;\n \n auto total_end = std::chrono::high_resolution_clock::now();\n auto total_duration = std::chrono::duration_cast(total_end - total_start);\n \n std::cout << \"Total Time: \" << cpu_duration.count() << \" ms\" << std::endl;\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/TensorParallelFromScratch/02_matmul_tp/matmul_tp_big.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n\n#include \"../../utils/utils.cuh\"\n\n#define M 2 // Rows of A, Rows of C\n#define N 2 // Columns of B, Columns of C\n#define K 3 // Columns of A, Rows of B\n\n// Kernel to perform our simple matrix multiplication\n__global__ void matMulKernel(int *A, int *B, int *C, int m, int n, int k) {\n int row = blockIdx.y * blockDim.y + threadIdx.y;\n int col = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"row: %d, col: %d\\n\", row, col);\n if (row < m && col < n) {\n int value = 0;\n for (int i = 0; i < k; i++) {\n value += A[row * k + i] * B[i * n + col];\n }\n C[row * n + col] = value;\n }\n}\n\nint main() {\n // Host matrices\n int A[M][K] = { {1, 2, 3}, {4, 5, 6} };\n int B[K][N] = { {100, 200}, {300, 400}, {500, 600} };\n int C[N][M] = {0};\n\n int *d_A, *d_B, *d_C;\n\n // Allocate device memory\n cudaCheckError(::cudaMalloc(&d_A, M * K * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_B, K * N * sizeof(int)));\n cudaCheckError(::cudaMalloc(&d_C, M * N * sizeof(int)));\n\n // Copy matrices from host to device\n cudaCheckError(::cudaMemcpy(d_A, A, M * K * sizeof(int), cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(d_B, B, K * N * sizeof(int), cudaMemcpyHostToDevice));\n\n // Set our grid and block size\n dim3 threadsPerBlock(M, N);\n dim3 numBlocks((N + threadsPerBlock.x - 1) / threadsPerBlock.x, (M + threadsPerBlock.y - 1) / threadsPerBlock.y);\n\n // Launch the kernel\n matMulKernel<<>>(d_A, d_B, d_C, M, N, K);\n\n // Copy result back to host\n cudaCheckError(::cudaMemcpy(C, d_C, M * N * sizeof(int), cudaMemcpyDeviceToHost));\n\n // Output the result matrix\n std::cout << \"Result Matrix C:\" << std::endl;\n for (int i = 0; i < M; i++) {\n for (int j = 0; j < N; j++) {\n std::cout << C[i][j] << \" \";\n }\n std::cout << std::endl;\n }\n\n // Free device memory\n cudaCheckError(::cudaFree(d_A));\n cudaCheckError(::cudaFree(d_B));\n cudaCheckError(::cudaFree(d_C));\n\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/TensorParallelFromScratch/01_simple_matmul_no_tp/matmul.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Simple matmul kernel.\n__global__ void matMul(float *out, const float *in1, const float *in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n float sum = 0.0f;\n if (row < width && col < width) {\n for (int i = 0; i < width; ++i) {\n sum += in1[row * width + i] * in2[i * width + col];\n }\n out[row * width + col] = sum;\n }\n}\n\n// Add one kernel to emphasize serial execution within a stream.\n__global__ void addOne(float *out, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n int idx = row * width + col;\n out[idx] += 1.0f;\n }\n}\n\n// Minus one kernel to emphasize serial execution within a stream.\n__global__ void minusOne(float *out, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n int idx = row * width + col;\n out[idx] -= 1.0f;\n }\n}\n\n// Add matrices to a final result for comparison.\n__global__ void matAdd(float *out, const float *in1, const float *in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n if (row < width && col < width) {\n out[row * width + col] = in1[row * width + col] + in2[row * width + col];\n }\n}\n\n// Perform a matmul on separate streams for a million elements, then add them.\nint main() {\n // Set dimensionality.\n const int width = 1024;\n const int size = width * width * sizeof(float);\n\n // Create separate streams to run matmuls asynchronously.\n cudaStream_t stream1, stream2;\n cudaCheckError(::cudaStreamCreate(&stream1));\n cudaCheckError(::cudaStreamCreate(&stream2));\n\n // Allocations and initializations.\n float *hA, *hB, *hC, *hD, *hResult;\n float *dA, *dB, *dC, *dD, *dResultAB, *dResultCD, *dFinal;\n\n // Host.\n cudaCheckError(::cudaMallocHost(&hA, size));\n cudaCheckError(::cudaMallocHost(&hB, size));\n cudaCheckError(::cudaMallocHost(&hC, size));\n cudaCheckError(::cudaMallocHost(&hD, size));\n cudaCheckError(::cudaMallocHost(&hResult, size));\n\n for (int i = 0; i < width * width; ++i) {\n hA[i] = 1.0f;\n hB[i] = 2.0f;\n hC[i] = 3.0f;\n hD[i] = 4.0f;\n }\n\n // Device.\n cudaCheckError(::cudaMalloc(&dA, size));\n cudaCheckError(::cudaMalloc(&dB, size));\n cudaCheckError(::cudaMalloc(&dC, size));\n cudaCheckError(::cudaMalloc(&dD, size));\n cudaCheckError(::cudaMalloc(&dResultAB, size));\n cudaCheckError(::cudaMalloc(&dResultCD, size));\n cudaCheckError(::cudaMalloc(&dFinal, size));\n\n // Copy data using separate streams. This is not necessary, per se, but it will\n // allow for overlap of copy and execution in a batch mode case.\n cudaCheckError(::cudaMemcpyAsync(dA, hA, size, cudaMemcpyHostToDevice, stream1));\n cudaCheckError(::cudaMemcpyAsync(dB, hB, size, cudaMemcpyHostToDevice, stream1));\n cudaCheckError(::cudaMemcpyAsync(dC, hC, size, cudaMemcpyHostToDevice, stream2));\n cudaCheckError(::cudaMemcpyAsync(dD, hD, size, cudaMemcpyHostToDevice, stream2));\n \n // Launch matmul kernels on different streams.\n dim3 threadsPerBlock(32, 32);\n dim3 numBlocks((width + threadsPerBlock.x - 1) / threadsPerBlock.x, \n (width + threadsPerBlock.y - 1) / threadsPerBlock.y);\n matMul<<>>(dResultAB, dA, dB, width);\n matMul<<>>(dResultCD, dC, dD, width);\n\n // Add 1.0f to each element in the result matrices\n addOne<<>>(dResultAB, width);\n addOne<<>>(dResultCD, width);\n\n // Subtract 1.0f to each element in the result matrices\n minusOne<<>>(dResultAB, width);\n minusOne<<>>(dResultCD, width);\n \n // Synchronize work prior to computing on each stream's result array.\n cudaCheckError(::cudaStreamSynchronize(stream1));\n cudaCheckError(::cudaStreamSynchronize(stream2));\n\n matAdd<<>>(dFinal, dResultAB, dResultCD, width);\n cudaCheckError(::cudaMemcpy(hResult, dFinal, size, cudaMemcpyDeviceToHost));\n\n // Verify results.\n float expected = (1.0f * 2.0f + 3.0f * 4.0f) * width;\n for (int i = 0; i < width * width; ++i) {\n assert(fabs(hResult[i] - expected) < 1e-5);\n }\n\n std::cout << \"Test passed!\" << std::endl;\n\n // Clean up.\n cudaCheckError(::cudaFree(dA));\n cudaCheckError(::cudaFree(dB));\n cudaCheckError(::cudaFree(dC));\n cudaCheckError(::cudaFree(dD));\n cudaCheckError(::cudaFree(dResultAB));\n cudaCheckError(::cudaFree(dResultCD));\n cudaCheckError(::cudaFree(dFinal));\n cudaCheckError(::cudaStreamDestroy(stream1));\n cudaCheckError(::cudaStreamDestroy(stream2));\n cudaCheckError(::cudaFreeHost(hA));\n cudaCheckError(::cudaFreeHost(hB));\n cudaCheckError(::cudaFreeHost(hC));\n cudaCheckError(::cudaFreeHost(hD));\n cudaCheckError(::cudaFreeHost(hResult));\n\n return 0;\n}\n\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/PerformanceChecklistExamples/cuda_streams/cuda_streams.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Device (GPU) code to print thread index from the GPU.\n__global__ void helloFromGPU()\n{\n int threadId = blockIdx.x * blockDim.x + threadIdx.x;\n printf(\"Hello from GPU thread: %d!\\n\", threadId);\n}\n\n// Host code\nint main()\n{\n int nDevices;\n cudaGetDeviceCount(&nDevices);\n for (int i = 0; i < nDevices; i++) {\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, i);\n printf(\"Device Number: %d\\n\", i);\n printf(\" Device name: %s\\n\", prop.name);\n printf(\" major-minor: %d-%d\\n\", prop.major, prop.minor);\n printf(\" Concurrent kernels: %s\\n\", prop.concurrentKernels ? \"yes\" : \"no\");\n printf(\" Concurrent computation/communication: %s\\n\\n\",prop.deviceOverlap ? \"yes\" : \"no\");\n printf(\" Managed memory supported: %s\\n\", prop.managedMemory ? \"yes\" : \"no\");\n printf(\" Memory Clock Rate (MHz): %d\\n\",\n prop.memoryClockRate/1024);\n printf(\" Memory Bus Width (bits): %d\\n\",\n prop.memoryBusWidth);\n printf(\" Peak Memory Bandwidth (GB/s): %.1f\\n\",\n 2.0*prop.memoryClockRate*(prop.memoryBusWidth/8)/1.0e6);\n printf(\" Total global memory (Gbytes): %.1f\\n\",(float)(prop.totalGlobalMem)/1024.0/1024.0/1024.0);\n printf(\" Shared memory per block (Kbytes): %.1f\\n\",(float)(prop.sharedMemPerBlock)/1024.0);\n printf(\" L2 cache size (Kbytes): %.1f\\n\",(float)(prop.l2CacheSize)/1024);\n printf(\" Warp-size: %d\\n\", prop.warpSize);\n }\n printf(\"\\n\");\n printf(\"Hello from CPU!\\n\");\n helloFromGPU<<<1,4>>>();\n cudaCheckError(::cudaDeviceSynchronize());\n return 0;\n}\n\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/SetupAndInitExamples/setup_check/hello_world.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define W 1024\n\n__global__ void matMul(float* out, const float* in1, const float* in2, const int width) {\n const int row = blockIdx.y * blockDim.y + threadIdx.y;\n const int col = blockIdx.x * blockDim.x + threadIdx.x;\n float sum = 0.0f;\n if (row < width && col < width) {\n for (int i = 0; i < width; ++i) {\n sum += in1[row * width + i] * in2[i * width + col];\n }\n out[row * width + col] = sum;\n }\n}\n\nvoid doMatMul(float* out, const float* in1, const float* in2, int width) {\n float *dIn1, *dIn2, *dOut;\n const size_t size = width * width * sizeof(float);\n\n printf(\"Allocating device memory...\\n\");\n cudaCheckError(::cudaMalloc(&dIn1, size));\n cudaCheckError(::cudaMalloc(&dIn2, size));\n cudaCheckError(::cudaMalloc(&dOut, size));\n\n printf(\"Copying data to device...\\n\");\n cudaCheckError(::cudaMemcpy(dIn1, in1, size, cudaMemcpyHostToDevice));\n cudaCheckError(::cudaMemcpy(dIn2, in2, size, cudaMemcpyHostToDevice));\n\n dim3 threadsPerBlock(16, 16);\n dim3 blocksPerGrid((width + threadsPerBlock.x - 1) / threadsPerBlock.x,\n (width + threadsPerBlock.y - 1) / threadsPerBlock.y);\n printf(\"Launching kernel...\\n\");\n for (int j = 0; j < 100; ++j)\n {\n nvtxRangePush(\"Matrix multiplication kernel\");\n matMul<<>>(dOut, dIn1, dIn2, width);\n nvtxRangePop();\n }\n cudaCheckError(::cudaDeviceSynchronize());\n\n printf(\"Copying result to host...\\n\");\n cudaCheckError(::cudaMemcpy(out, dOut, size, cudaMemcpyHostToDevice));\n\n printf(\"Freeing device memory...\\n\");\n cudaCheckError(::cudaFree(dIn1));\n cudaCheckError(::cudaFree(dIn2));\n cudaCheckError(::cudaFree(dOut));\n}\n\nint main() {\n float *in1, *in2, *out;\n const size_t width = W * W * sizeof(float);\n cudaCheckError(::cudaMallocHost(&in1, width));\n cudaCheckError(::cudaMallocHost(&in2, width));\n cudaCheckError(::cudaMallocHost(&out, width));\n for (int i = 0; i < W * W; ++i) {\n in1[i] = static_cast(rand()) / RAND_MAX;\n in2[i] = static_cast(rand()) / RAND_MAX;\n }\n printf(\"Starting matrix multiplication...\\n\");\n doMatMul(out, in1, in2, width);\n printf(\"Matrix multiplication completed.\\n\");\n printf(\"Freeing host memory...\\n\");\n cudaCheckError(::cudaFreeHost(in1));\n cudaCheckError(::cudaFreeHost(in2));\n cudaCheckError(::cudaFreeHost(out));\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/ProfilingExamples/nvtx/nvtx_example.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define N (1 << 20) // 1M elements\n#define THREADS 1024 // Threads per block\n#define BLOCKS (N + THREADS - 1) / THREADS // Number of blocks\n\n__global__ void reduce_global(float *input, float *output) {\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x * 2 + tid; // Unroll by 2 for better memory efficiency\n\n float sum = 0.0f;\n if (global_idx < N) sum += input[global_idx];\n if (global_idx + blockDim.x < N) sum += input[global_idx + blockDim.x];\n\n // In-register reduction (warp-level)\n for (int stride = warpSize / 2; stride > 0; stride >>= 1) {\n sum += __shfl_down_sync(0xFFFFFFFF, sum, stride);\n }\n\n // Store only the first thread's result per warp\n if (tid % warpSize == 0) {\n atomicAdd(&output[blockIdx.x], sum);\n }\n}\n\n\n\n// Optimized Reduction with Shared Memory\n__global__ void reduce_shared(float *input, float *output) {\n __shared__ float shared_mem[THREADS + (THREADS / 32)];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load data into shared memory (coalesced)\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// CPU Reduction (Baseline for Benchmark)\nfloat reduce_cpu(float *input, int size) {\n float sum = 0.0f;\n for (int i = 0; i < size; i++) {\n sum += input[i];\n }\n return sum;\n}\n\n// Benchmarking Function\nvoid benchmark() {\n float *h_input, *h_output;\n float *d_input, *d_output;\n cudaCheckError(::cudaMallocHost(&h_input, N * sizeof(float)));\n cudaCheckError(::cudaMallocHost(&h_output, BLOCKS * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_input, N * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_output, BLOCKS * sizeof(float)));\n\n // Initialize input\n for (int i = 0; i < N; i++) {\n h_input[i] = 1.0f; // Simple case where sum = N\n }\n\n cudaCheckError(::cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));\n\n cudaEvent_t start, stop;\n cudaCheckError(::cudaEventCreate((&start)));\n cudaCheckError(::cudaEventCreate(&stop));\n // Benchmark Global Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_global<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_naive;\n cudaCheckError(::cudaEventElapsedTime(&time_naive, start, stop));\n printf(\"Global Kernel Time: %f ms\\n\", time_naive);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n\n // Final sum on CPU\n float sum_naive = reduce_cpu(h_output, BLOCKS);\n\n // Benchmark Optimized Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_shared<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_optimized;\n cudaCheckError(::cudaEventElapsedTime(&time_optimized, start, stop));\n printf(\"Shmem Kernel Time: %f ms\\n\", time_optimized);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n \n float sum_optimized = reduce_cpu(h_output, BLOCKS);\n\n // CPU Baseline\n float sum_cpu = reduce_cpu(h_input, N);\n\n printf(\"CPU Sum: %f\\n\", sum_cpu);\n printf(\"GPU Global memory Sum: %f\\n\", sum_naive);\n printf(\"GPU Shared memory Sum: %f\\n\", sum_optimized);\n\n // Cleanup\n cudaCheckError(::cudaFreeHost(h_input));\n cudaCheckError(::cudaFreeHost(h_output));\n cudaCheckError(::cudaFree(d_input));\n cudaCheckError(::cudaFree(d_output));\n}\n\nint main() {\n benchmark();\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/ProfilingExamples/global_vs_shared_mem/global_vs_shared.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n#define N (1 << 20) // 1M elements\n#define THREADS 1024 // Threads per block\n#define BLOCKS (N + THREADS - 1) / THREADS // Number of blocks\n\n// Naive Reduction (with bank conflicts)\n__global__ void reduce_naive(float *input, float *output) {\n __shared__ float shared_mem[THREADS];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load global memory into shared memory\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n // Reduction in shared memory\n for (int stride = 1; stride < blockDim.x; stride *= 2) {\n if (tid % (2 * stride) == 0) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// Optimized Reduction (Avoiding Bank Conflicts)\n__global__ void reduce_optimized(float *input, float *output) {\n __shared__ float shared_mem[THREADS + (THREADS / 32)];\n\n int tid = threadIdx.x;\n int global_idx = blockIdx.x * blockDim.x + tid;\n\n // Load data into shared memory (coalesced)\n shared_mem[tid] = (global_idx < N) ? input[global_idx] : 0.0f;\n __syncthreads();\n\n for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n shared_mem[tid] += shared_mem[tid + stride];\n }\n __syncthreads();\n }\n\n // Write block sum to output\n if (tid == 0) output[blockIdx.x] = shared_mem[0];\n}\n\n// CPU Reduction (Baseline for Benchmark)\nfloat reduce_cpu(float *input, int size) {\n float sum = 0.0f;\n for (int i = 0; i < size; i++) {\n sum += input[i];\n }\n return sum;\n}\n\n// Benchmarking Function\nvoid benchmark() {\n float *h_input, *h_output;\n float *d_input, *d_output;\n cudaCheckError(::cudaMallocHost(&h_input, N * sizeof(float)));\n cudaCheckError(::cudaMallocHost(&h_output, BLOCKS * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_input, N * sizeof(float)));\n cudaCheckError(::cudaMalloc(&d_output, BLOCKS * sizeof(float)));\n\n // Initialize input\n for (int i = 0; i < N; i++) {\n h_input[i] = 1.0f; // Simple case where sum = N\n }\n\n cudaCheckError(::cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));\n\n cudaEvent_t start, stop;\n cudaCheckError(::cudaEventCreate((&start)));\n cudaCheckError(::cudaEventCreate(&stop));\n // Benchmark Na\u00efve Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_naive<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_naive;\n cudaCheckError(::cudaEventElapsedTime(&time_naive, start, stop));\n printf(\"Na\u00efve Kernel Time: %f ms\\n\", time_naive);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n\n // Final sum on CPU\n float sum_naive = reduce_cpu(h_output, BLOCKS);\n\n // Benchmark Optimized Kernel\n cudaCheckError(::cudaEventRecord(start));\n reduce_optimized<<>>(d_input, d_output);\n cudaCheckError(::cudaEventRecord(stop));\n cudaCheckError(::cudaEventSynchronize(stop));\n float time_optimized;\n cudaCheckError(::cudaEventElapsedTime(&time_optimized, start, stop));\n printf(\"Optimized Kernel Time: %f ms\\n\", time_optimized);\n cudaCheckError(::cudaMemcpy(h_output, d_output, BLOCKS * sizeof(float), cudaMemcpyDeviceToHost));\n \n float sum_optimized = reduce_cpu(h_output, BLOCKS);\n\n // CPU Baseline\n float sum_cpu = reduce_cpu(h_input, N);\n\n printf(\"CPU Sum: %f\\n\", sum_cpu);\n printf(\"Na\u00efve GPU Sum: %f\\n\", sum_naive);\n printf(\"Optimized GPU Sum: %f\\n\", sum_optimized);\n\n // Cleanup\n cudaCheckError(::cudaFreeHost(h_input));\n cudaCheckError(::cudaFreeHost(h_output));\n cudaCheckError(::cudaFree(d_input));\n cudaCheckError(::cudaFree(d_output));\n}\n\nint main() {\n benchmark();\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/ProfilingExamples/memory_bank_conflict/bank_conflict.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Simple kernel for use.\n__global__ void vector_add_kernel(float *d_v1, float *d_v2, float *d_v3, int n) {\n int i = blockDim.x * blockIdx.x + threadIdx.x;\n if (i < n) {\n d_v3[i] = d_v1[i] + d_v2[i];\n }\n}\n\n// Structure to pass around host memory pointers.\nstruct Memory {\n float *v1;\n float *v2;\n float *v3;\n};\n\nvoid benchmark(int N, bool usePinnedMemory) {\n size_t size = N * sizeof(float);\n Memory mem;\n \n // Allocate host-pinned memory.\n if (usePinnedMemory) {\n cudaCheckError(::cudaMallocHost(&mem.v1, size));\n cudaCheckError(::cudaMallocHost(&mem.v2, size));\n cudaCheckError(::cudaMallocHost(&mem.v3, size));\n } else {\n // Allocate pageable memory.\n mem.v1 = (float*)malloc(size);\n mem.v2 = (float*)malloc(size);\n mem.v3 = (float*)malloc(size);\n }\n\n // Initialize data\n for (int i = 0; i < N; i++) {\n mem.v1[i] = 1.0f;\n mem.v2[i] = 2.0f;\n }\n\n // Allocate device memory\n float *d_v1, *d_v2, *d_v3;\n cudaCheckError(::cudaMalloc(&d_v1, size));\n cudaCheckError(::cudaMalloc(&d_v2, size));\n cudaCheckError(::cudaMalloc(&d_v3, size));\n\n // Create CUDA stream\n cudaStream_t stream;\n cudaCheckError(::cudaStreamCreate(&stream));\n\n // Create CUDA events for timing\n cudaEvent_t start, stop;\n cudaCheckError(::cudaEventCreate(&start));\n cudaCheckError(::cudaEventCreate(&stop));\n\n // Record start time\n cudaCheckError(::cudaEventRecord(start, stream));\n\n // Asynchronous memory transfers\n cudaCheckError(::cudaMemcpyAsync(d_v1, mem.v1, size, cudaMemcpyHostToDevice, stream));\n cudaCheckError(::cudaMemcpyAsync(d_v2, mem.v2, size, cudaMemcpyHostToDevice, stream));\n\n // Launch kernel\n int threadsPerBlock = 256;\n int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;\n vector_add_kernel<<>>(d_v1, d_v2, d_v3, N);\n\n // Asynchronous memory transfer back to host\n cudaCheckError(::cudaMemcpyAsync(mem.v3, d_v3, size, cudaMemcpyDeviceToHost, stream));\n\n // Record stop time\n cudaCheckError(::cudaEventRecord(stop, stream));\n cudaCheckError(::cudaEventSynchronize(stop));\n\n float milliseconds = 0;\n cudaCheckError(::cudaEventElapsedTime(&milliseconds, start, stop));\n\n // Calculate bandwidth\n float bandwidth = 3 * size / (milliseconds / 1000) / 1e9; // GB/s\n\n std::cout << (usePinnedMemory ? \"Pinned\" : \"Pageable\") << \" memory, N = \" << N \n << \", Time: \" << milliseconds << \" ms, Bandwidth: \" << bandwidth << \" GB/s\" << std::endl;\n\n // Clean up\n cudaCheckError(::cudaFree(d_v1));\n cudaCheckError(::cudaFree(d_v2));\n cudaCheckError(::cudaFree(d_v3));\n cudaCheckError(::cudaStreamDestroy(stream));\n cudaCheckError(::cudaEventDestroy(start));\n cudaCheckError(::cudaEventDestroy(stop));\n\n if (usePinnedMemory) {\n cudaCheckError(::cudaFreeHost(mem.v1));\n cudaCheckError(::cudaFreeHost(mem.v2));\n cudaCheckError(::cudaFreeHost(mem.v3));\n } else {\n free(mem.v1);\n free(mem.v2);\n free(mem.v3);\n }\n}\n\nint main() {\n // Small, to fairly large memory pushes:\n // 1024 * 4 bytes = 4KB\n // 1024 * 1024 * 4 bytes = 4MB\n // 10 * 1024 * 1024 * 4 bytes = 40MB\n std::vector sizes = {1024, 1024 * 1024, 10 * 1024 * 1024};\n \n for (int N : sizes) {\n benchmark(N, false); // Pageable memory\n benchmark(N, true); // Pinned memory\n std::cout << std::endl;\n }\n\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/MemoryAndStructureExamples/pinned_vs_pageable/pinned_vs_pageable.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"../../utils/utils.cuh\"\n\n// Kernel to intialize weights with random values.\n__global__ void initializeWeightsKernel(float* weights, int size, unsigned long seed) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < size) {\n curandStateXORWOW state;\n curand_init(seed, idx, 0, &state);\n weights[idx] = curand_normal(&state);\n }\n}\n\n// Function to initialize weights using cudaMalloc\nvoid initializeWeightsCudaMalloc(int size) {\n\n float* d_weights;\n // Allocate memory for distribution, but don't initialize values until kernel\n cudaCheckError(::cudaMalloc(&d_weights, size * sizeof(float)));\n\n // Launch kernel to initialize weights\n initializeWeightsKernel<<<(size + 255) / 256, 256>>>(d_weights, size, time(NULL));\n cudaCheckError(::cudaDeviceSynchronize());\n\n cudaCheckError(::cudaFree(d_weights));\n\n std::cout << \"Initialized weights using cudaMalloc.\" << std::endl;\n}\n\n// Function to initialize weights using thrust::device_vector\nvoid initializeWeightsThrust(int size) {\n\n // Create a device vector and initialize with zeros\n thrust::device_vector weights(size);\n\n // Launch kernel to initialize weights\n initializeWeightsKernel<<<(size + 255) / 256, 256>>>(thrust::raw_pointer_cast(weights.data()), size, time(NULL));\n\n std::cout << \"Initialized weights using thrust::device_vector.\" << std::endl;\n}\n\nint main() {\n const int size = 1000000; // Large number of weights\n\n // Initialize weights using thrust::device_vector\n initializeWeightsThrust(size);\n\n // Initialize weights using cudaMalloc\n initializeWeightsCudaMalloc(size);\n\n return 0;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/cuda_examples/MemoryAndStructureExamples/alloc_init_vs_alloc_uninit/alloc_init_vs_alloc_uninit.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \nnamespace py = pybind11;\n\n#define gpuErrchk(ans) \\\n { \\\n gpuAssert((ans), __FILE__, __LINE__); \\\n }\ninline void gpuAssert(cudaError_t code, const char* file, int line, bool abort = true) {\n if (code != cudaSuccess) {\n fprintf(stderr, \"GPUassert %d: %s %s %d\\n\", (int)code, cudaGetErrorString(code), file, line);\n if (abort) exit(code);\n }\n}\n\n// find the magnitude after enforcing the 2:4 sparsity constraint on a group of 4 values\n__device__ float group_2_to_4(float4 vals) {\n vals.x = fabs(vals.x);\n vals.y = fabs(vals.y);\n vals.z = fabs(vals.z);\n vals.w = fabs(vals.w);\n\n float sum0 = vals.x + vals.y;\n float sum1 = vals.x + vals.z;\n float sum2 = vals.x + vals.w;\n float sum3 = vals.y + vals.z;\n float sum4 = vals.y + vals.w;\n float sum5 = vals.z + vals.w;\n\n float best_sum0 = fmax(sum0, sum1);\n float best_sum1 = fmax(sum2, sum3);\n float best_sum2 = fmax(sum4, sum5);\n float best_sum = fmax(fmax(best_sum0, best_sum1), best_sum2);\n\n return best_sum;\n}\n\ninline float* float_ptr_from_numpy(py::array_t& py_float) { return (float*)py_float.data(); }\n\ninline unsigned int* uint_ptr_from_numpy(py::array_t& py_uint) { return (unsigned int*)py_uint.data(); }\n\n/**********************************************************\n * Check for the best permutation for an entire matrix\n **********************************************************/\n__global__ void permute_and_sum_after_2_to_4(float* matrix, unsigned int rows, unsigned int cols, unsigned int* stripes,\n unsigned int total_stripes, unsigned int* permutations, float* output) {\n // vectorize\n float4* mat4 = (float4*)matrix;\n cols /= 4;\n\n // each thread in a block takes some number of rows\n size_t num_rows = max((int)ceilf((float)rows / (float)blockDim.x), 1);\n size_t row_offset = num_rows * threadIdx.x;\n size_t num_stripes = total_stripes; // total_stripes / gridDim.x;\n size_t stripe_offset = 0; // num_stripes * blockIdx.x;\n unsigned int localStart = stripe_offset;\n unsigned int localEnd = localStart + num_stripes;\n\n // each block takes care of one permutation\n unsigned int p = blockIdx.x;\n unsigned int* permutation = &permutations[p * total_stripes * 4];\n\n float sum = 0.0f;\n extern __shared__ float s[32][32];\n float4* local_stripes = (float4*)&s[threadIdx.x];\n float* local_columns = (float*)&s[threadIdx.x];\n float4* permuted_local_stripes = (float4*)&local_stripes[num_stripes];\n float* permuted_local_columns = (float*)&local_columns[num_stripes * 4];\n\n for (unsigned int r = row_offset; r < row_offset + num_rows; ++r) {\n if (r >= rows) break;\n\n // load into smem\n for (unsigned int s = localStart; s < localEnd; ++s) {\n unsigned int stripe = stripes[s];\n local_stripes[s] = mat4[r * cols + stripe];\n }\n\n// now permute\n#pragma unroll 4\n for (unsigned int c = 0; c < num_stripes * 4; ++c) {\n permuted_local_columns[c] = local_columns[permutation[c]];\n }\n\n // now sum 2:4\n for (unsigned int s = 0; s < num_stripes; ++s) {\n sum += group_2_to_4(permuted_local_stripes[s]);\n }\n }\n\n atomicAdd(&output[p], sum);\n}\n\nvoid free_permutation_memory(float** dmatrix, unsigned int** dstripe_groups, unsigned int** dpermutations,\n float** dresults, float** hresults) {\n cudaFree(*dmatrix);\n cudaFree(*dresults);\n cudaFree(*dpermutations);\n cudaFree(*dstripe_groups);\n free(*hresults);\n}\n\nint set_up_check_permutation_memory(float** dmatrix, unsigned int rows, unsigned int cols,\n unsigned int** dstripe_groups, unsigned int group_width, unsigned int num_groups,\n unsigned int** dpermutations, unsigned int num_permutations, float** dresults,\n float** hresults) {\n static unsigned int setupRows = 0;\n static unsigned int setupCols = 0;\n static unsigned int setupGroupWidth = 0;\n static unsigned int setupNumGroups = 0;\n static unsigned int setupNumPermutations = 0;\n static bool allocated = false;\n int fresh_alloc = 0;\n if (!allocated || setupRows != rows || setupCols != cols || setupGroupWidth != group_width ||\n setupNumGroups != num_groups || setupNumPermutations != num_permutations) {\n if (allocated) {\n free_permutation_memory(dmatrix, dstripe_groups, dpermutations, dresults, hresults);\n }\n\n gpuErrchk(cudaMalloc((void**)dmatrix, rows * cols * sizeof(float)));\n gpuErrchk(cudaMalloc((void**)dstripe_groups, group_width * num_groups * sizeof(unsigned int)));\n gpuErrchk(cudaMalloc((void**)dpermutations, num_permutations * group_width * 4 * sizeof(unsigned int)));\n gpuErrchk(cudaMalloc((void**)dresults, num_permutations * sizeof(float)));\n *hresults = (float*)malloc(num_permutations * sizeof(float));\n\n allocated = true;\n setupRows = rows;\n setupCols = cols;\n setupGroupWidth = group_width;\n setupNumGroups = num_groups;\n setupNumPermutations = num_permutations;\n fresh_alloc = 1;\n }\n\n return fresh_alloc;\n}\n\nint run_check_permutations(\n py::array_t& py_matrix, unsigned int rows, unsigned int cols,\n py::array_t&\n py_stripe_groups, // groups of stripes, group_width = stripes per group, num_groups = groups in the array\n unsigned int group_width, unsigned int num_groups,\n py::array_t& py_permutations, // array of permutations to try, group_width*4 values per each of\n // num_permutations permutations\n unsigned int num_permutations,\n py::array_t& py_improvement, // improvment offered by the best permutation\n py::array_t& py_permutation // the best permutation\n) {\n const unsigned int threads = 32;\n static float* d_matrix;\n static unsigned int* d_permutations;\n static unsigned int* d_stripes;\n static float* d_results;\n static float* results;\n\n float* matrix = float_ptr_from_numpy(py_matrix);\n unsigned int* stripe_groups = uint_ptr_from_numpy(py_stripe_groups);\n unsigned int* permutations = uint_ptr_from_numpy(py_permutations);\n float* improvement = float_ptr_from_numpy(py_improvement);\n unsigned int* permutation = uint_ptr_from_numpy(py_permutation);\n\n int fresh_alloc = set_up_check_permutation_memory(&d_matrix, rows, cols, &d_stripes, group_width, num_groups,\n &d_permutations, num_permutations, &d_results, &results);\n if (fresh_alloc == 1) {\n gpuErrchk(cudaMemcpy(d_permutations, permutations, num_permutations * group_width * 4 * sizeof(unsigned int),\n cudaMemcpyHostToDevice));\n gpuErrchk(\n cudaMemcpy(d_stripes, stripe_groups, group_width * num_groups * sizeof(unsigned int), cudaMemcpyHostToDevice));\n }\n\n // initialize results, new matrix\n gpuErrchk(cudaMemset(d_results, 0, num_permutations * sizeof(float)));\n gpuErrchk(cudaMemcpy(d_matrix, matrix, rows * cols * sizeof(float), cudaMemcpyHostToDevice));\n\n // get results for all permutations\n permute_and_sum_after_2_to_4<<>>(\n d_matrix, rows, cols, d_stripes, group_width, d_permutations, d_results);\n gpuErrchk(cudaDeviceSynchronize());\n\n gpuErrchk(cudaMemcpy(results, d_results, num_permutations * sizeof(float), cudaMemcpyDeviceToHost));\n\n // find the best permutation - could reduce on GPU\n unsigned int best_permutation = 0;\n float best_improvement = 0.0f;\n for (unsigned int p = 1; p < num_permutations; ++p) {\n float cur_improvement = results[p] - results[0];\n if (best_improvement < cur_improvement) {\n best_permutation = p;\n best_improvement = cur_improvement;\n }\n }\n\n *improvement = best_improvement;\n *permutation = best_permutation;\n\n return 0;\n}\n\n///////////////////////////////////////////////////////////\n\n/**********************************************************\n * Get the magnitude of a matrix after applying 2:4\n **********************************************************/\n// find the magnitude after enforcing the 2:4 sparsity constraint on a subset of the columns of an input matrix\n__global__ void subset_sum_after_2_to_4(float* matrix, unsigned int rows, unsigned int cols, unsigned int start_col,\n unsigned int end_col, float* output) {\n // vectorize\n float4* mat4 = (float4*)matrix;\n cols /= 4;\n start_col /= 4;\n end_col /= 4;\n\n // each thread in a block takes some number of rows\n size_t num_rows = max((int)ceilf((float)rows / (float)blockDim.x), 1);\n size_t row_offset = num_rows * threadIdx.x;\n // each block takes some number of columns\n size_t num_cols = (end_col - start_col) / gridDim.x;\n size_t col_offset = num_cols * blockIdx.x;\n start_col += col_offset;\n end_col = start_col + num_cols;\n\n float sum = 0.0f;\n for (unsigned int r = row_offset; r < row_offset + num_rows; ++r) {\n if (r < rows) {\n for (unsigned int c = start_col; c < end_col; c++) {\n sum += group_2_to_4(mat4[r * cols + c]);\n }\n }\n }\n\n atomicAdd(output, sum);\n}\n\n// build the entire permute map at once\n// each block handles one group of stripes\n// each threads in the block handle all handle the same permutation at the same time on different rows before moving to\n// the next permutation\n__global__ void build_permute_map(float* matrix, unsigned int rows, unsigned int cols, unsigned int* stripes,\n unsigned int group_width, unsigned int* permutations, unsigned int num_permutations,\n unsigned int perm_length, float* output, unsigned int* best_indices) {\n // vectorize\n float4* mat4 = (float4*)matrix;\n cols /= 4;\n\n // each block handles a group of stripes\n unsigned int* stripe_group = (unsigned int*)&stripes[blockIdx.x * group_width];\n\n // shared memory: 32 threads each need 16*2\n extern __shared__ float pm_shared[32][32];\n float4* local_stripes = (float4*)&pm_shared[threadIdx.x];\n float* local_columns = (float*)&pm_shared[threadIdx.x];\n float4* permuted_stripes = (float4*)&local_stripes[4];\n float* permuted_columns = (float*)&local_columns[16];\n\n // each thread handles all permutations in the row before moving on to the next row\n size_t num_rows = max((int)ceilf((float)rows / (float)blockDim.x), 1);\n size_t row_offset = num_rows * threadIdx.x;\n\n for (unsigned int r = row_offset; r < row_offset + num_rows; ++r) {\n if (r >= rows) break;\n\n // load a row into smem\n for (unsigned int s = 0; s < group_width; ++s) {\n unsigned int const stripe = stripe_group[s];\n local_stripes[s] = mat4[r * cols + stripe];\n }\n\n for (unsigned int p = 0; p < num_permutations; ++p) {\n unsigned int* permutation = &permutations[p * perm_length];\n float sum = 0.0f;\n\n// permute\n#pragma unroll 4\n for (unsigned int c = 0; c < group_width * 4; ++c) {\n permuted_columns[c] = local_columns[permutation[c]];\n }\n\n // sum 2:4\n for (unsigned int s = 0; s < group_width; ++s) {\n sum += group_2_to_4(permuted_stripes[s]);\n }\n\n // update the running sum for this stripe group's permutation\n atomicAdd(&output[blockIdx.x * num_permutations + p], sum);\n }\n }\n\n // at this point, each permutation's sum in this stripe group has been calculated\n // now, find the best option\n __syncthreads();\n\n if (threadIdx.x == 0) {\n unsigned int best_permutation = 0;\n float best_magnitude = output[blockIdx.x * num_permutations];\n float base_magnitude = best_magnitude;\n\n // #pragma unroll 32\n for (unsigned int p = 1; p < num_permutations; ++p) {\n float magnitude = output[blockIdx.x * num_permutations + p];\n if (magnitude > best_magnitude) {\n best_permutation = p;\n best_magnitude = magnitude;\n }\n }\n\n output[blockIdx.x * num_permutations] = best_magnitude - base_magnitude;\n best_indices[blockIdx.x] = best_permutation;\n }\n}\n\nvoid free_sum_after_2_to_4_memory(float** dmatrix, float** dresult) {\n cudaFree(*dmatrix);\n cudaFree(*dresult);\n}\n\nint set_up_sum_after_2_to_4_memory(float** dmatrix, unsigned int rows, unsigned int cols, float** dresult) {\n static unsigned int setupRows = 0;\n static unsigned int setupCols = 0;\n static bool allocated = false;\n\n int fresh_allocation = 0;\n if (!allocated || setupRows != rows || setupCols != cols) {\n if (allocated) free_sum_after_2_to_4_memory(dmatrix, dresult);\n\n gpuErrchk(cudaMalloc((void**)dmatrix, rows * cols * sizeof(float)));\n gpuErrchk(cudaMalloc((void**)dresult, sizeof(float)));\n\n setupRows = rows;\n setupCols = cols;\n\n fresh_allocation = 1;\n }\n\n allocated = true;\n\n return fresh_allocation;\n}\n\nint run_subset_sum_after_2_to_4(py::array_t& py_matrix, unsigned int rows, unsigned int cols,\n unsigned int start_col, unsigned int end_col, unsigned int blocks, unsigned int threads,\n py::array_t& py_output) {\n static float* d_matrix;\n static float* d_result;\n\n int fresh_allocation = set_up_sum_after_2_to_4_memory(&d_matrix, rows, cols, &d_result);\n\n float* matrix = float_ptr_from_numpy(py_matrix);\n float* output = float_ptr_from_numpy(py_output);\n\n gpuErrchk(cudaMemcpy(d_matrix, matrix, rows * cols * sizeof(float), cudaMemcpyHostToDevice));\n gpuErrchk(cudaMemset(d_result, 0, sizeof(float)));\n\n subset_sum_after_2_to_4<<>>(d_matrix, rows, cols, start_col, end_col, d_result);\n gpuErrchk(cudaDeviceSynchronize());\n\n gpuErrchk(cudaMemcpy(output, d_result, sizeof(float), cudaMemcpyDeviceToHost));\n\n return 0;\n}\n\nvoid set_up_permute_map_memory(float** dmatrix, unsigned int rows, unsigned int cols, unsigned int** dstripes,\n unsigned int num_groups, unsigned int group_width, unsigned int** dpermutations,\n unsigned int num_permutations, unsigned int perm_length, float** doutput,\n unsigned int** dindices, float** hresult, unsigned int** hindices) {\n static unsigned int setUpRows = 0;\n static unsigned int setUpCols = 0;\n static unsigned int setUpGroupWidth = 0;\n static unsigned int setUpNumGroups = 0;\n static unsigned int setUpNumPerms = 0;\n static unsigned int setUpPermLength = 0;\n\n if (setUpRows != rows || setUpCols != cols) {\n if (*dmatrix != NULL) {\n gpuErrchk(cudaFree(*dmatrix));\n *dmatrix = NULL;\n }\n gpuErrchk(cudaMalloc((void**)dmatrix, rows * cols * sizeof(float)));\n }\n\n if (setUpGroupWidth < group_width || setUpNumGroups < num_groups) {\n if (*dstripes != NULL) {\n gpuErrchk(cudaFree(*dstripes));\n *dstripes = NULL;\n }\n gpuErrchk(cudaMalloc((void**)dstripes, num_groups * group_width * sizeof(unsigned int)));\n\n if (setUpNumGroups < num_groups) {\n if (*dindices != NULL) {\n gpuErrchk(cudaFree(*dindices));\n *dindices = NULL;\n }\n gpuErrchk(cudaMalloc((void**)dindices, num_groups * sizeof(unsigned int)));\n if (*hindices != NULL) {\n free(*hindices);\n *hindices = NULL;\n }\n *hindices = (unsigned int*)malloc(num_groups * sizeof(unsigned int));\n }\n }\n\n if (setUpNumPerms < num_permutations || setUpPermLength < perm_length) {\n if (*dpermutations != NULL) {\n gpuErrchk(cudaFree(*dpermutations));\n *dpermutations = NULL;\n }\n gpuErrchk(cudaMalloc((void**)dpermutations, perm_length * num_permutations * sizeof(unsigned int)));\n }\n\n if (setUpNumPerms < num_permutations || setUpNumGroups < num_groups) {\n if (*doutput != NULL) {\n gpuErrchk(cudaFree(*doutput));\n *doutput = NULL;\n }\n gpuErrchk(cudaMalloc((void**)doutput, num_permutations * num_groups * sizeof(float)));\n if (*hresult != NULL) {\n free(*hresult);\n *hresult = NULL;\n }\n *hresult = (float*)malloc(num_permutations * num_groups * sizeof(float));\n }\n\n setUpRows = rows;\n setUpCols = cols;\n setUpGroupWidth = group_width;\n setUpNumGroups = num_groups;\n setUpNumPerms = num_permutations;\n setUpPermLength = perm_length;\n}\n\nint run_build_permute_map(py::array_t& py_matrix, unsigned int rows, unsigned int cols,\n py::array_t& py_stripes, unsigned int num_groups, unsigned int group_width,\n py::array_t& py_permutations, unsigned int perm_length,\n py::array_t& py_improvements, py::array_t& py_best_indices) {\n static float* d_matrix = NULL;\n static unsigned int* d_stripes = NULL;\n static unsigned int* d_permutations = NULL;\n static float* d_output = NULL;\n static unsigned int* d_indices = NULL;\n static float* hresult = NULL;\n static unsigned int* hindices = NULL;\n\n const unsigned int num_permutations = py_permutations.size() / perm_length;\n\n const unsigned int MAX_GROUPS_PER_LAUNCH = num_permutations <= 5775 ? 1820 : 40;\n const unsigned int full_launches = num_groups / MAX_GROUPS_PER_LAUNCH;\n const unsigned int final_launch = num_groups % MAX_GROUPS_PER_LAUNCH;\n const unsigned int launches = full_launches + (final_launch != 0 ? 1 : 0);\n\n set_up_permute_map_memory(&d_matrix, rows, cols, &d_stripes, min(num_groups, MAX_GROUPS_PER_LAUNCH), group_width,\n &d_permutations, num_permutations, perm_length, &d_output, &d_indices, &hresult, &hindices);\n\n float* matrix = float_ptr_from_numpy(py_matrix);\n unsigned int* stripes = uint_ptr_from_numpy(py_stripes);\n unsigned int* permutations = uint_ptr_from_numpy(py_permutations);\n float* improvements = float_ptr_from_numpy(py_improvements);\n unsigned int* best_indices = uint_ptr_from_numpy(py_best_indices);\n\n gpuErrchk(cudaMemcpy(d_matrix, matrix, rows * cols * sizeof(float), cudaMemcpyHostToDevice));\n gpuErrchk(cudaMemcpy(d_permutations, permutations, num_permutations * perm_length * sizeof(unsigned int),\n cudaMemcpyHostToDevice));\n\n unsigned int group_offset = 0;\n for (unsigned int l = 0; l < launches; ++l) {\n unsigned int groups_this_launch = (l < full_launches) ? MAX_GROUPS_PER_LAUNCH : final_launch;\n\n gpuErrchk(cudaMemcpy(d_stripes, &stripes[group_offset * group_width],\n groups_this_launch * group_width * sizeof(unsigned int), cudaMemcpyHostToDevice));\n gpuErrchk(cudaMemset(d_output, 0, groups_this_launch * num_permutations * sizeof(float)));\n gpuErrchk(cudaMemset(d_indices, 0, groups_this_launch * sizeof(unsigned int)));\n\n unsigned int shmem = 32 * (32) * sizeof(float);\n build_permute_map<<>>(d_matrix, rows, cols, d_stripes, group_width, d_permutations,\n num_permutations, perm_length, d_output, d_indices);\n gpuErrchk(cudaDeviceSynchronize());\n\n gpuErrchk(\n cudaMemcpy(hresult, d_output, num_permutations * groups_this_launch * sizeof(float), cudaMemcpyDeviceToHost));\n gpuErrchk(cudaMemcpy(hindices, d_indices, groups_this_launch * sizeof(unsigned int), cudaMemcpyDeviceToHost));\n\n // thread0 stuck the minimum in the first slot of each group\n for (unsigned int g = 0; g < groups_this_launch; ++g) {\n improvements[group_offset + g] = hresult[g * num_permutations];\n best_indices[group_offset + g] = hindices[g];\n }\n\n group_offset += groups_this_launch;\n }\n\n return 0;\n}\n\n/**********************************************************\n * Build the swap map for channel_swaps\n **********************************************************/\n// find the magnitude improvement if some columns were swapped (check all pairs of columns in all the stripe_pairs)\n__global__ void swap_columns_sum_after_2_to_4(float* matrix, unsigned int rows, unsigned int cols,\n unsigned int* stripe_pairs, float* output) {\n // vectorize\n float4* mat4 = (float4*)matrix;\n cols /= 4;\n\n // each thread takes some number of rows\n size_t const num_rows = max((int)ceilf((float)rows / (float)blockDim.x), 1);\n size_t const row_offset = num_rows * threadIdx.x;\n\n // each block is repsonsible for a pair of stripes\n unsigned int const stripe0 = stripe_pairs[2 * blockIdx.x];\n unsigned int const stripe1 = stripe_pairs[2 * blockIdx.x + 1];\n // space for 32 threads, 8 values (2 stripes) in use at a time, plus 16 partial sums and one base sum\n extern __shared__ float cs[32][32];\n float4* local_stripe0 = (float4*)&cs[threadIdx.x][0];\n float* local_cols0 = (float*)&cs[threadIdx.x][0];\n float4* local_stripe1 = (float4*)&cs[threadIdx.x][4];\n float* local_cols1 = (float*)&cs[threadIdx.x][4];\n float* local_psum = (float*)&cs[threadIdx.x][8];\n float* base_psum = (float*)&cs[threadIdx.x][24];\n\n *base_psum = 0.0f;\n for (unsigned int s = 0; s < 16; ++s) {\n local_psum[s] = 0.0f;\n }\n\n for (unsigned int r = row_offset; r < row_offset + num_rows; ++r) {\n if (r >= rows) break;\n *local_stripe0 = mat4[r * cols + stripe0];\n *local_stripe1 = mat4[r * cols + stripe1];\n *base_psum += group_2_to_4(*local_stripe0) + group_2_to_4(*local_stripe1);\n unsigned int swap_idx = 0;\n for (unsigned int c0 = 0; c0 < 4; ++c0) {\n for (unsigned int c1 = 0; c1 < 4; ++c1) {\n // swap c0 and c1\n float tmp = local_cols0[c0];\n local_cols0[c0] = local_cols1[c1];\n local_cols1[c1] = tmp;\n\n // grab the sum\n local_psum[swap_idx] += group_2_to_4(*local_stripe0) + group_2_to_4(*local_stripe1);\n\n // swap back\n local_cols1[c1] = local_cols0[c0];\n local_cols0[c0] = tmp;\n\n swap_idx++;\n }\n }\n }\n\n // reduce partial sums, store local diffs in the output\n __syncthreads();\n if (threadIdx.x == 0) {\n for (unsigned int t = 1; t < blockDim.x; ++t) {\n for (unsigned int swap = 0; swap < 16; ++swap) {\n local_psum[swap] += cs[t][8 + swap];\n }\n *base_psum += cs[t][24];\n }\n\n for (unsigned int swap = 0; swap < 16; ++swap) {\n atomicAdd(&output[blockIdx.x * 16 + swap], local_psum[swap] - (*base_psum));\n }\n }\n}\n\nvoid set_up_swap_map_memory(float** dmatrix, unsigned int rows, unsigned int cols, unsigned int** dstripe_pairs,\n unsigned int num_pairs, float** dresult) {\n static unsigned int setupRows = 0;\n static unsigned int setupCols = 0;\n static unsigned int setupPairs = 0;\n\n if (*dmatrix == NULL || setupRows != rows || setupCols != cols) {\n if (*dmatrix != NULL) {\n gpuErrchk(cudaFree(*dmatrix));\n *dmatrix = NULL;\n }\n gpuErrchk(cudaMalloc((void**)dmatrix, rows * cols * sizeof(float)));\n setupRows = rows;\n setupCols = cols;\n }\n\n if (*dstripe_pairs == NULL || *dresult == NULL || setupPairs < num_pairs) {\n if (*dstripe_pairs != NULL) {\n gpuErrchk(cudaFree(*dstripe_pairs));\n *dstripe_pairs = NULL;\n }\n if (*dresult != NULL) {\n gpuErrchk(cudaFree(*dresult));\n *dresult = NULL;\n }\n gpuErrchk(cudaMalloc((void**)dstripe_pairs, num_pairs * 2 * sizeof(unsigned int)));\n gpuErrchk(cudaMalloc((void**)dresult, num_pairs * 16 * sizeof(float)));\n\n setupPairs = num_pairs;\n }\n}\n\nint run_build_swap_map(py::array_t& py_matrix, unsigned int rows, unsigned int cols,\n py::array_t& py_stripe_pairs, py::array_t& py_output) {\n static float* d_matrix = NULL;\n static float* d_result = NULL;\n static unsigned int* d_stripe_pairs = NULL;\n\n float* matrix = float_ptr_from_numpy(py_matrix); //(float*)py_matrix.data();\n unsigned int* stripe_pairs = uint_ptr_from_numpy(py_stripe_pairs); //(unsigned int*)py_stripe_pairs.data();\n float* output = float_ptr_from_numpy(py_output); //(float*)py_output.data();\n\n unsigned int num_pairs = py_stripe_pairs.size() / 2;\n\n set_up_swap_map_memory(&d_matrix, rows, cols, &d_stripe_pairs, num_pairs, &d_result);\n gpuErrchk(cudaMemcpy(d_matrix, matrix, rows * cols * sizeof(float), cudaMemcpyHostToDevice));\n gpuErrchk(cudaMemcpy(d_stripe_pairs, stripe_pairs, 2 * num_pairs * sizeof(unsigned int), cudaMemcpyHostToDevice));\n gpuErrchk(cudaMemset(d_result, 0, num_pairs * 16 * sizeof(float)));\n\n unsigned int shmem = 32 * (32) * sizeof(float);\n swap_columns_sum_after_2_to_4<<>>(d_matrix, rows, cols, d_stripe_pairs, d_result);\n gpuErrchk(cudaDeviceSynchronize());\n\n gpuErrchk(cudaMemcpy(output, d_result, num_pairs * 16 * sizeof(float), cudaMemcpyDeviceToHost));\n\n return 0;\n}\n///////////////////////////////////////////////////////////\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"sum_after_2_to_4\", &run_subset_sum_after_2_to_4, \"matrix sum after applying 2:4 (CUDA)\",\n py::call_guard());\n m.def(\"build_permute_map\", &run_build_permute_map, \"optimize stripe groups (CUDA)\",\n py::call_guard());\n m.def(\"check_permutations\", &run_check_permutations, \"exhaustively check all permutations (CUDA)\",\n py::call_guard());\n m.def(\"build_swap_map\", &run_build_swap_map, \"channel swaps (CUDA)\", py::call_guard());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/sparsity/permutation_search_kernels/CUDA_kernels/permutation_search_kernels.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<128, 64, 16, 1, 4, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_128_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\nvoid run_fmha_dgrad_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 128 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_128_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_128_64_sm80_kernel<<>>(params);\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_128_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<512, 64, 16, 1, 8, 0x00u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_512_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\ntemplate \n__global__ void fmha_fprop_fp16_512_64_sm80_kernel_nl(Fused_multihead_attention_fprop_params params,\n const int num_full_heads, const int num_main_groups,\n const int main_group_size, const int main_steps,\n const int rest_steps) {\n fmha::device_1xN(params, num_full_heads, num_main_groups, main_group_size, main_steps,\n rest_steps);\n}\n\nvoid run_fmha_fp16_512_64_sm80_(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_512_64_sm80_kernel\n : &fmha_fprop_fp16_512_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\nvoid run_fmha_fp16_512_64_sm80_nl_(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_512_64_sm80_kernel_nl\n : &fmha_fprop_fp16_512_64_sm80_kernel_nl;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n if (configure) {\n const int heads_total = launch_params.params.b * launch_params.params.h;\n std::tie(launch_params.num_full_heads, launch_params.num_main_groups, launch_params.heads_last_wave,\n launch_params.main_steps, launch_params.rest_steps, launch_params.elts_per_thread) =\n fmha::work_dist(total_ctas, heads_total);\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(\n launch_params.params, launch_params.num_full_heads, launch_params.num_main_groups, launch_params.heads_last_wave,\n launch_params.main_steps, launch_params.rest_steps);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\nvoid run_fmha_fp16_512_64_sm80(Launch_params& launch_params,\n const bool configure) {\n if (launch_params.is_nl) {\n run_fmha_fp16_512_64_sm80_nl_(launch_params, configure);\n } else {\n run_fmha_fp16_512_64_sm80_(launch_params, configure);\n }\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_512_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<128, 64, 16, 1, 4, 0x08u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_128_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\nvoid run_fmha_fp16_128_64_sm80(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_128_64_sm80_kernel\n : &fmha_fprop_fp16_128_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_128_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<256, 64, 16, 1, 4, 0x08u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_256_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\nvoid run_fmha_fp16_256_64_sm80(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_256_64_sm80_kernel\n : &fmha_fprop_fp16_256_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_256_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n#include \"fmha_dgrad_kernel_1xN_reload_nl.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<512, 64, 16, 1, 8, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_512_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\ntemplate \n__global__ void fmha_dgrad_fp16_512_64_sm80_nl_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN_nl(params);\n fmha::compute_dq_dk_1xN_nl(params);\n}\n\nvoid run_fmha_dgrad_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 512 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_512_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_512_64_sm80_kernel<<>>(params);\n}\n\nvoid run_fmha_dgrad_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params& params, const int num_chunks,\n cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 512 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n auto kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<2>;\n\n if (num_chunks == 2) {\n kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<2>;\n } else if (num_chunks == 3) {\n kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<3>;\n } else {\n assert(false && \"Unsupperted number of chunks\");\n }\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n dim3 grid(params.h, params.b, num_chunks);\n\n kernel<<>>(params);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_512_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n\ninline __device__ float4 ldg128(const void* ptr) { return *static_cast(ptr); }\n\ninline __device__ void stg128(void* ptr, const float4& data) { *static_cast(ptr) = data; }\n\ntemplate \n__global__ __launch_bounds__(THREADS) void fmha_noloop_reduce_kernel(void* __restrict__ out,\n const void* __restrict__ in,\n const int* __restrict__ cu_seqlens,\n const int batch_size) {\n enum { BYTES_PER_LDG = 16 };\n enum { NUM_ELTS = BYTES_PER_LDG / sizeof(T) };\n\n // One CTA hidden vector for K and V\n enum { BYTES_PER_ROW = HIDDEN_SIZE * sizeof(T) * 2 };\n // The stride in bytes in dQKV\n enum { OUT_STRIDE_BYTES = 3 * HIDDEN_SIZE * sizeof(T) };\n // The offset in bytes in dQKV to the dKV part for non-interleaved heads\n enum { OUT_OFFSET_KV_BYTES = HIDDEN_SIZE * sizeof(T) };\n\n static_assert(BYTES_PER_ROW == HIDDEN_SIZE * 2 * sizeof(T));\n\n // Size in bytes of the input tile\n enum { BYTES_PER_TILE = CHUNKS * BYTES_PER_ROW };\n\n enum { BYTES_PER_CTA = THREADS * BYTES_PER_LDG };\n\n enum { LDGS = BYTES_PER_ROW / BYTES_PER_CTA };\n static_assert(BYTES_PER_CTA * LDGS == BYTES_PER_ROW);\n\n union Vec_t {\n float4 raw;\n T elt[NUM_ELTS];\n };\n\n // ZERO-OUT invalid positions in dQKV\n const int total = cu_seqlens[batch_size];\n if (blockIdx.x >= total) {\n enum { BYTES_PER_QKV_ROW = 3 * HIDDEN_SIZE * sizeof(T) };\n enum { STGS = BYTES_PER_QKV_ROW / BYTES_PER_LDG };\n\n const float4 zeros = make_float4(0.f, 0.f, 0.f, 0.f);\n\n char* base_ptr = static_cast(out) + blockIdx.x * OUT_STRIDE_BYTES;\n\n for (int tidx = threadIdx.x; tidx < STGS; tidx += THREADS) {\n stg128(base_ptr + tidx * BYTES_PER_LDG, zeros);\n }\n\n return;\n }\n\n // SETUP\n const int offset_in = blockIdx.x * BYTES_PER_TILE + threadIdx.x * BYTES_PER_LDG;\n const char* ptr_in = static_cast(in) + offset_in;\n\n const int offset_out = blockIdx.x * OUT_STRIDE_BYTES + threadIdx.x * BYTES_PER_LDG;\n char* ptr_out = static_cast(out) + OUT_OFFSET_KV_BYTES + offset_out;\n\n // LOAD\n\n Vec_t local_in[CHUNKS][LDGS];\n\n#pragma unroll\n for (int c = 0; c < CHUNKS; c++) {\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n int offset = c * BYTES_PER_ROW + l * BYTES_PER_CTA;\n local_in[c][l].raw = ldg128(ptr_in + offset);\n }\n }\n\n // UNPACK\n float acc[LDGS][NUM_ELTS];\n\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n#pragma unroll\n for (int e = 0; e < NUM_ELTS; e++) {\n acc[l][e] = float(local_in[0][l].elt[e]);\n }\n }\n\n// COMPUTE\n#pragma unroll\n for (int c = 1; c < CHUNKS; c++) {\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n#pragma unroll\n for (int e = 0; e < NUM_ELTS; e++) {\n acc[l][e] += float(local_in[c][l].elt[e]);\n }\n }\n }\n\n // PACK\n Vec_t local_out[LDGS];\n\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n#pragma unroll\n for (int e = 0; e < NUM_ELTS; e++) {\n local_out[l].elt[e] = T(acc[l][e]);\n }\n }\n\n// STORE\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n const int offset = l * BYTES_PER_CTA;\n stg128(ptr_out + offset, local_out[l].raw);\n }\n}\n\nvoid fmha_run_noloop_reduce(void* out, const void* in, const int* cu_seqlens, const int hidden_size,\n const int batch_size, const int total, const int num_chunks, cudaStream_t stream) {\n const int blocks = total;\n\n if (hidden_size == 1024) {\n constexpr int HIDDEN_SIZE = 1024;\n constexpr int THREADS = 256;\n\n if (num_chunks == 2) {\n fmha_noloop_reduce_kernel\n <<>>(out, in, cu_seqlens, batch_size);\n } else if (num_chunks == 3) {\n fmha_noloop_reduce_kernel\n <<>>(out, in, cu_seqlens, batch_size);\n } else {\n assert(false && \"Unsupported num_chunks\");\n }\n\n } else {\n assert(false && \"Unsupported hidden_size\");\n }\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_noloop_reduce.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<384, 64, 16, 1, 4, 0x18u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_384_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\nvoid run_fmha_fp16_384_64_sm80(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_384_64_sm80_kernel\n : &fmha_fprop_fp16_384_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_384_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<384, 64, 16, 1, 8, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_384_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\nvoid run_fmha_dgrad_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 384 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_384_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_384_64_sm80_kernel<<>>(params);\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_384_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<256, 64, 16, 1, 4, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_256_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\nvoid run_fmha_dgrad_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 256 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_256_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_256_64_sm80_kernel<<>>(params);\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_256_64_kernel.sm80.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n\nconstexpr int block_size = 512;\nconstexpr int ctas_per_sm = 4;\n\ntemplate \n__global__ void __launch_bounds__(block_size)\n mha_fill_kernel(scalar_t* out_tensor, const int32_t* const start_row, const size_t num_rows) {\n size_t row_stride = gridDim.y * blockDim.x;\n size_t row_index = blockIdx.x + (size_t)start_row[0];\n size_t col_index = blockIdx.y * blockDim.x + threadIdx.x;\n while (row_index < num_rows) {\n out_tensor[row_index * row_stride + col_index] = 0;\n row_index += gridDim.x;\n }\n}\n\nat::Tensor& mha_fill(at::Tensor& self, const at::Tensor& start_index) {\n auto max_tokens = self.size(0);\n auto self_2d = self.view({max_tokens, -1});\n auto fcd_size = self_2d.size(1);\n TORCH_CHECK(self.is_contiguous(), \"input not contiguous\");\n TORCH_CHECK(fcd_size % block_size == 0, \"input size not aligned to block size\");\n const int num_mp = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n uint64_t num_blk_y = (uint64_t)(fcd_size / block_size);\n uint64_t num_blk_x = (uint64_t)std::ceil(num_mp * ctas_per_sm / num_blk_y);\n dim3 dim_grid(num_blk_x, num_blk_y);\n dim3 dim_block(block_size);\n\n AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(\n at::ScalarType::Half, at::ScalarType::BFloat16, self_2d.scalar_type(), \"mha_padding_fill_\", [&]() {\n mha_fill_kernel<<>>(\n self_2d.data_ptr(), start_index.data_ptr(), max_tokens);\n C10_CUDA_KERNEL_LAUNCH_CHECK();\n });\n return self;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/fmha/src/fmha_fill.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n\n// Use 128-bit vectorization\ntypedef uint4 vector_t;\n\n#define ASSERT_ALIGNED(DTYPE, PTR) \\\n TORCH_INTERNAL_ASSERT(is_aligned(PTR), \"Tensor \" #PTR \" is not \" #DTYPE \" aligned\")\n\ntemplate \nbool is_aligned(const void* ptr) noexcept {\n auto iptr = reinterpret_cast(ptr);\n return !(iptr % alignof(T));\n}\n\ntemplate \n__global__ void focal_loss_forward_cuda_kernel(outscalar_t* loss, scalar_t* partial_grad,\n const scalar_t* __restrict__ cls_output,\n const labelscalar_t* __restrict__ cls_targets_at_level,\n const float* __restrict__ num_positives_sum, const int64_t num_examples,\n const int64_t num_classes, const int64_t num_real_classes,\n const float alpha, const float gamma, const float smoothing_factor) {\n extern __shared__ unsigned char shm[];\n accscalar_t* loss_shm = reinterpret_cast(shm);\n loss_shm[threadIdx.x] = 0;\n accscalar_t loss_acc = 0;\n\n accscalar_t one = accscalar_t(1.0);\n accscalar_t K = accscalar_t(2.0);\n accscalar_t normalizer = one / static_cast(num_positives_sum[0]);\n accscalar_t nn_norm, np_norm, pn_norm, pp_norm;\n\n // *_norm is used for label smoothing only\n if (SMOOTHING) {\n nn_norm = one - smoothing_factor / K;\n np_norm = smoothing_factor / K;\n pn_norm = smoothing_factor - smoothing_factor / K;\n pp_norm = one - smoothing_factor + smoothing_factor / K;\n }\n\n vector_t p_vec, grad_vec;\n\n // Accumulate loss on each thread\n for (int64_t i = (blockIdx.x * blockDim.x + threadIdx.x) * ILP; i < num_examples * num_classes;\n i += gridDim.x * blockDim.x * ILP) {\n int64_t idy = i / num_classes;\n labelscalar_t y = cls_targets_at_level[idy];\n int64_t base_yid = i % num_classes;\n\n int64_t pos_idx = idy * num_classes + y;\n p_vec = *(vector_t*)&cls_output[i]; // Vectorized load\n\n // Skip ignored matches\n if (y == -2) {\n#pragma unroll\n for (int j = 0; j < ILP; j++) {\n *((scalar_t*)(&grad_vec) + j) = 0;\n }\n *(vector_t*)&partial_grad[i] = grad_vec;\n continue;\n }\n\n#pragma unroll\n for (int j = 0; j < ILP; j++) {\n // Skip the pad classes\n if (base_yid + j >= num_real_classes) {\n *((scalar_t*)(&grad_vec) + j) = 0;\n continue;\n }\n\n accscalar_t p = static_cast(*((scalar_t*)(&p_vec) + j));\n accscalar_t exp_np = ::exp(-p);\n accscalar_t exp_pp = ::exp(p);\n accscalar_t sigma = one / (one + exp_np);\n accscalar_t logee = (p >= 0) ? exp_np : exp_pp;\n accscalar_t addee = (p >= 0) ? 0 : -p;\n accscalar_t off_a = addee + ::log(one + logee);\n\n // Negative matches\n accscalar_t base = SMOOTHING ? nn_norm * p : p;\n accscalar_t off_b = (SMOOTHING ? np_norm : 0) - sigma;\n accscalar_t coeff_f1 = one - alpha;\n accscalar_t coeff_f2 = sigma;\n accscalar_t coeff_b1 = gamma;\n accscalar_t coeff_b2 = one - sigma;\n\n // Positive matches\n if (y >= 0 && (i + j == pos_idx)) {\n base = SMOOTHING ? pn_norm * p : 0;\n off_b = (SMOOTHING ? pp_norm : one) - sigma;\n coeff_f1 = alpha;\n coeff_f2 = one - sigma;\n coeff_b1 = -gamma;\n coeff_b2 = sigma;\n }\n\n accscalar_t coeff_f = coeff_f1 * ::pow(coeff_f2, gamma);\n accscalar_t coeff_b = coeff_b1 * coeff_b2;\n\n accscalar_t loss_t = coeff_f * (base + off_a);\n accscalar_t grad = coeff_f * (coeff_b * (base + off_a) - off_b);\n\n // Delay the normalize of partial gradient by num_positives_sum to back\n // propagation because scalar_t reduces precision. Focal loss is very\n // sensitive to the small gradient. No worry on overflow here since\n // gradient has relative smaller range than input.\n loss_acc += loss_t;\n *((scalar_t*)(&grad_vec) + j) = static_cast(grad);\n }\n\n // This may generate two vectorized stores instead of one\n *(vector_t*)&partial_grad[i] = grad_vec;\n }\n loss_shm[threadIdx.x] = loss_acc;\n\n // Intra-CTA reduction\n __syncthreads();\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n loss_shm[threadIdx.x] += loss_shm[threadIdx.x + s];\n }\n __syncthreads();\n }\n\n // Inter-CTA reduction\n if (threadIdx.x == 0) {\n loss_acc = loss_shm[0] * normalizer;\n atomicAdd(loss, loss_acc);\n }\n}\n\ntemplate \n__global__ void focal_loss_backward_cuda_kernel(scalar_t* partial_grad, const outscalar_t* __restrict__ grad_output,\n const float* __restrict__ num_positives_sum, const uint64_t numel) {\n int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * ILP;\n\n accscalar_t normalizer = static_cast(grad_output[0]) / static_cast(num_positives_sum[0]);\n\n // The input is enforced to pad to use vector load, thus there's no need to\n // check whether the last element of ILP can out of bound.\n if (idx >= numel) return;\n\n vector_t grad_vec;\n grad_vec = *(vector_t*)&partial_grad[idx];\n#pragma unroll(ILP)\n for (int i = 0; i < ILP; i++) {\n auto grad = static_cast(*((scalar_t*)(&grad_vec) + i));\n grad *= normalizer;\n *((scalar_t*)(&grad_vec) + i) = static_cast(grad);\n }\n *(vector_t*)&partial_grad[idx] = grad_vec;\n}\n\nstd::vector focal_loss_forward_cuda(const at::Tensor& cls_output, const at::Tensor& cls_targets_at_level,\n const at::Tensor& num_positives_sum, const int64_t num_real_classes,\n const float alpha, const float gamma, const float smoothing_factor) {\n // Checks required for correctness\n TORCH_INTERNAL_ASSERT(cls_output.size(-1) >= num_real_classes, \"Incorrect number of real classes.\");\n TORCH_INTERNAL_ASSERT(cls_targets_at_level.scalar_type() == at::kLong, \"Invalid label type.\");\n TORCH_INTERNAL_ASSERT((num_positives_sum.numel() == 1) && (num_positives_sum.scalar_type() == at::kFloat),\n \"Expect num_positives_sum to be a float32 tensor with only one element.\");\n TORCH_INTERNAL_ASSERT(cls_output.dim() == cls_targets_at_level.dim() + 1,\n \"Mis-matched dimensions between class output and label.\");\n for (int64_t i = 0; i < cls_targets_at_level.dim(); i++)\n TORCH_INTERNAL_ASSERT(cls_output.size(i) == cls_targets_at_level.size(i),\n \"Mis-matched shape between class output and label.\");\n\n // Checks required for better performance\n const int ILP = sizeof(vector_t) / cls_output.element_size();\n ASSERT_ALIGNED(vector_t, cls_output.data_ptr());\n TORCH_INTERNAL_ASSERT(cls_output.size(-1) % ILP == 0,\n \"Pad number of classes first to take advantage of vectorized load.\");\n TORCH_INTERNAL_ASSERT(num_real_classes >= ILP, \"Too few classes.\");\n\n int64_t num_classes = cls_output.size(-1);\n int64_t num_examples = cls_output.numel() / num_classes;\n at::Tensor loss = at::zeros({}, cls_output.options().dtype(at::kFloat));\n\n // Compute the incompelete gradient during fprop since most of the heavy\n // functions of bprop are the same as fprop, thus trade memory for compute\n // helps with focal loss.\n at::Tensor partial_grad = at::empty_like(cls_output);\n\n // Set the number of CTAs per SM according to the compute capability.\n // Each CTA loops on input with stride till the last item.\n cudaDeviceProp props;\n cudaGetDeviceProperties(&props, at::cuda::current_device());\n int cta_per_sm = 2;\n if (props.major >= 10) {\n cta_per_sm = 8;\n }\n dim3 block(512);\n dim3 grid(cta_per_sm * props.multiProcessorCount);\n\n // Specialize on label smoothing or not to reduce redundant operations\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (smoothing_factor == 0.0f) {\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(cls_output.scalar_type(), \"focal_loss_fprop\", [&] {\n using accscalar_t = at::acc_type;\n using labelscalar_t = int64_t;\n using outscalar_t = float;\n const int ILP = sizeof(vector_t) / sizeof(scalar_t);\n focal_loss_forward_cuda_kernel\n <<>>(\n loss.data_ptr(), partial_grad.data_ptr(), cls_output.data_ptr(),\n cls_targets_at_level.data_ptr(), num_positives_sum.data_ptr(), num_examples,\n num_classes, num_real_classes, alpha, gamma, smoothing_factor);\n });\n } else {\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(cls_output.scalar_type(), \"focal_loss_fprop\", [&] {\n using accscalar_t = at::acc_type;\n using labelscalar_t = int64_t;\n using outscalar_t = float;\n const int ILP = sizeof(vector_t) / sizeof(scalar_t);\n focal_loss_forward_cuda_kernel\n <<>>(\n loss.data_ptr(), partial_grad.data_ptr(), cls_output.data_ptr(),\n cls_targets_at_level.data_ptr(), num_positives_sum.data_ptr(), num_examples,\n num_classes, num_real_classes, alpha, gamma, smoothing_factor);\n });\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n return {loss, partial_grad};\n}\n\nat::Tensor focal_loss_backward_cuda(const at::Tensor& grad_output, const at::Tensor& partial_grad,\n const at::Tensor& num_positives_sum) {\n // Each thread process ILP elements\n const int ILP = sizeof(vector_t) / partial_grad.element_size();\n dim3 block(512);\n dim3 grid((partial_grad.numel() + block.x * ILP - 1) / (block.x * ILP));\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(partial_grad.scalar_type(), \"focal_loss_bprop\", [&] {\n using accscalar_t = at::acc_type;\n using outscalar_t = float;\n const int ILP = sizeof(vector_t) / sizeof(scalar_t);\n focal_loss_backward_cuda_kernel\n <<>>(partial_grad.data_ptr(), grad_output.data_ptr(),\n num_positives_sum.data_ptr(), partial_grad.numel());\n });\n\n AT_CUDA_CHECK(cudaGetLastError());\n return partial_grad;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/focal_loss/focal_loss_cuda_kernel.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n\n#include \n\n#include \"ATen/ATen.h\"\n#include \"ATen/TensorUtils.h\"\n#include \"ATen/cuda/CUDAContext.h\"\n#include \"ATen/cuda/detail/IndexUtils.cuh\"\n// #include \"ATen/Type.h\"\n#include \"ATen/AccumulateType.h\"\n#include \"multi_tensor_apply.cuh\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\n#include \"type_shim.h\"\n\ntypedef enum {\n ADAM_MODE_0 = 0, // eps under square root\n ADAM_MODE_1 = 1 // eps outside square root\n} adamMode_t;\n\ntemplate \n__global__ void adam_cuda_kernel(T* __restrict__ p,\n GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale, const float step_size,\n const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n for (int j = i; j < tsize; j += totThreads) {\n T scaled_grad = g[j] / grad_scale;\n m[j] = b1 * m[j] + (1 - b1) * scaled_grad;\n v[j] = b2 * v[j] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * p[j]);\n p[j] = p[j] - (step_size * update);\n if (p_copy != NULL) p_copy[j] = (GRAD_T)p[j];\n }\n}\n\ntemplate \nstruct AdamFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata& tl,\n const float b1, const float b2, const float eps, const float grad_scale,\n const float step_size, adamMode_t mode, const float decay) {\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n T* p = (T*)tl.addresses[0][tensor_loc];\n p += chunk_idx * chunk_size;\n T* m = (T*)tl.addresses[1][tensor_loc];\n m += chunk_idx * chunk_size;\n T* v = (T*)tl.addresses[2][tensor_loc];\n v += chunk_idx * chunk_size;\n GRAD_T* g = (GRAD_T*)tl.addresses[3][tensor_loc];\n g += chunk_idx * chunk_size;\n GRAD_T* p_copy = NULL;\n if (DEPTH == 5) {\n p_copy = (GRAD_T*)tl.addresses[4][tensor_loc];\n p_copy += chunk_idx * chunk_size;\n }\n\n n -= chunk_idx * chunk_size;\n\n T incoming_p[ILP];\n T incoming_m[ILP];\n T incoming_v[ILP];\n T incoming_g[ILP];\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(p) && is_aligned(m) && is_aligned(v) && is_aligned(g) &&\n is_aligned(p_copy)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n GRAD_T tmp_g[ILP];\n load_store(incoming_p, p, 0, i_start);\n load_store(incoming_m, m, 0, i_start);\n load_store(incoming_v, v, 0, i_start);\n load_store(tmp_g, g, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_g[ii] = static_cast(tmp_g[ii]);\n T scaled_grad = incoming_g[ii] / grad_scale;\n incoming_m[ii] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n incoming_v[ii] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(incoming_v[ii] + eps);\n else // Mode 1\n denom = sqrtf(incoming_v[ii]) + eps;\n float update = (incoming_m[ii] / denom) + (decay * incoming_p[ii]);\n incoming_p[ii] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) tmp_g[ii] = static_cast(incoming_p[ii]);\n }\n load_store(p, incoming_p, i_start, 0);\n load_store(m, incoming_m, i_start, 0);\n load_store(v, incoming_v, i_start, 0);\n if (DEPTH == 5) load_store(p_copy, tmp_g, i_start, 0);\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_p[ii] = 0;\n incoming_m[ii] = 0;\n incoming_v[ii] = 0;\n incoming_g[ii] = 0;\n\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n incoming_p[ii] = p[i];\n incoming_m[ii] = m[i];\n incoming_v[ii] = v[i];\n incoming_g[ii] = static_cast(g[i]);\n }\n }\n\n // note for clarification to future michael:\n // From a pure memory dependency perspective, there's likely no point unrolling\n // the write loop, since writes just fire off once their LDGs arrive.\n // Put another way, the STGs are dependent on the LDGs, but not on each other.\n // There is still compute ILP benefit from unrolling the loop though.\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = i_start + threadIdx.x + ii * blockDim.x;\n\n if (j < n && j < chunk_size) {\n T scaled_grad = incoming_g[ii] / grad_scale;\n m[j] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n v[j] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * incoming_p[ii]);\n p[j] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) p_copy[j] = (GRAD_T)p[j];\n }\n }\n }\n }\n }\n};\n\nvoid fused_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr,\n float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction,\n float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n adam_cuda_kernel<<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_adam_cuda_mt(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, // p, m, v, g, p_copy\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 4 || tl_sz == 5, \"expected tensor lists of size 4 or 5\");\n\n if (tensor_lists[3][0].scalar_type() == at::ScalarType::Half) {\n // alher values should be fp32 for half gradients\n TORCH_CHECK(tensor_lists[0][0].scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dich is done on the gradient type\n if (tl_sz == 5) {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n } else {\n if (tl_sz == 5) {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\ntemplate \n__device__ void convert(const FROM_T vi, TO_T& vo) {\n vo = static_cast(vi);\n}\n\ntemplate <>\n__device__ void convert(const float vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = vi;\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, float& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = static_cast(t.as_half);\n}\n\ntemplate <>\n__device__ void convert(const at::Half vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = static_cast(vi);\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, at::Half& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = t.as_half;\n}\n\ntemplate \n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, GRAD_T* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n GRAD_T pi = p_copy[j];\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\ntemplate <>\n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, uint8_t* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n at::Half pi;\n convert(p_copy[j], pi);\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\n\ntemplate \n__global__ void maybe_cast_kernel(volatile int* overflow_flag, const FROM_T* p_in, TO_T* p_out, const size_t tsize) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = 0;\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n p_out[j] = po[ii];\n }\n }\n }\n}\n\ntemplate \n__global__ void reversible_adam_cuda_kernel(\n T* __restrict__ p,\n REDU_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1, const float b2, const float eps,\n const float grad_scale, const float step_size, const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n bool overflow = false;\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n mi[ii] = b1 * mi[ii] + (1 - b1) * scaled_grad;\n vi[ii] = b2 * vi[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n float update = (mi[ii] / denom) + (decay * pi[ii]);\n pi[ii] = pi[ii] - (step_size * update);\n } else {\n overflow = true;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n if (p_copy != NULL) {\n convert(pi[ii], p_copy[j]);\n }\n }\n }\n }\n\n if (p_copy != NULL) {\n __syncthreads();\n if (overflow) {\n convert(float(INFINITY), p_copy[0]);\n }\n }\n}\n\ntemplate \n__global__ void maybe_adam_undo_cuda_kernel(volatile int* overflow_flag, T* __restrict__ p, T* __restrict__ m,\n T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale,\n const float step_size, const size_t tsize, adamMode_t mode,\n const float decay) {\n // NB! Skip undo kernel when overflow flag is NOT set\n if (overflow_flag && *overflow_flag == 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i * ILP;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n pi[ii] = (pi[ii] + step_size * (mi[ii] / denom)) / (1.0f - step_size * decay);\n mi[ii] = (mi[ii] - (1 - b1) * scaled_grad) / b1;\n vi[ii] = (vi[ii] - (1 - b2) * scaled_grad * scaled_grad) / b2;\n // Make sure round off errors don't create (small) negative value.\n // This can happen if we have to revert the very first step.\n vi[ii] = vi[ii] >= 0.0f ? vi[ii] : 0.0f;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i * ILP;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n }\n }\n }\n}\n\ntemplate \nstruct MaybeCastFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* overflow_flag,\n TensorListMetadata& tl) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n FROM_T* p_in = (FROM_T*)tl.addresses[0][tensor_loc];\n p_in += chunk_idx * chunk_size;\n TO_T* p_out = (TO_T*)tl.addresses[1][tensor_loc];\n p_out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n int dim = chunk_size < n ? chunk_size : n;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < dim; j_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = FROM_T(0);\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n p_out[j] = po[ii];\n }\n }\n }\n }\n};\n\nvoid fused_strided_check_finite(at::Tensor& overflow_flag, at::Tensor& p_copy, int stride, int clear_overflow_first) {\n // Get tensor size\n int tsize = p_copy.numel();\n int niter = (tsize + stride - 1) / stride;\n\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n // In order to avoid race condition, blocks must be 1 when clear_overflow_first flag is set.\n const dim3 blocks(clear_overflow_first ? 1 : (niter + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_copy), \"parameter tensor is too large to be indexed with int32\");\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_HALF_AND_BYTE(\n p_copy.scalar_type(), 0, \"check_finite_cuda_kernel\",\n strided_check_finite_cuda_kernel<<>>(\n overflow_flag.data_ptr(), p_copy.data_ptr(), tsize, stride, clear_overflow_first););\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_reversible_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n if (p_copy.numel() == 0 || p_copy.scalar_type() == g.scalar_type()) {\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps,\n grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n TORCH_CHECK(p_copy.scalar_type() == at::ScalarType::Byte, \"expected parameter to be of byte type\");\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_e5m2_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.data_ptr(), m.data_ptr(),\n v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale, step_size, tsize,\n (adamMode_t)mode, decay););\n }\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n reversible_adam_cuda_kernel\n <<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid maybe_cast_cuda(at::Tensor& overflow_flag, at::Tensor& p_in, at::Tensor& p_out) {\n // Get tensor size\n int tsize = p_in.numel();\n TORCH_CHECK(tsize == p_out.numel(), \"p_in.numel() must equal p_out.numel()\");\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_in), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n DISPATCH_FLOAT_HALF_AND_BYTE(p_in.scalar_type(), 0,\n \"maybe_cast_cuda\" DISPATCH_FLOAT_HALF_AND_BYTE(\n p_out.scalar_type(), 1, \"maybe_cast_cuda\",\n maybe_cast_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL,\n p_in.data_ptr(), p_out.data_ptr(), tsize);))\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid maybe_cast_cuda_mt(int chunk_size, at::Tensor overflow_flag,\n std::vector> tensor_lists) // p_in, p_out\n{\n // Constants\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 2, \"expected tensor lists of size 2\");\n\n DISPATCH_FLOAT_HALF_AND_BYTE(\n tensor_lists[0][0].scalar_type(), 0, \"maybe_cast_cuda_mt_kernel\",\n DISPATCH_FLOAT_HALF_AND_BYTE(tensor_lists[1][0].scalar_type(), 1, \"maybe_cast_cuda_mt_kernel\",\n multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, overflow_flag, tensor_lists,\n MaybeCastFunctor<2, scalar_t_0, scalar_t_1>());))\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_maybe_adam_undo_cuda(at::Tensor& overflow_flag, at::Tensor& p, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n maybe_adam_undo_cuda_kernel\n <<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(\n g.scalar_type(), 0, \"adam_cuda_kernel\",\n maybe_adam_undo_cuda_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale,\n step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n#define CUDACHECK(cmd) \\\n do { \\\n cudaError_t err = cmd; \\\n if (err != cudaSuccess) { \\\n char hostname[1024]; \\\n gethostname(hostname, 1024); \\\n printf(\"%s: CUDA failure %s:%d '%s'\\n\", hostname, __FILE__, __LINE__, cudaGetErrorString(err)); \\\n } \\\n } while (0)\n\nnamespace {\n\nconstexpr int THREADS_PER_CTA = 128;\n\n/* Basic deleter function for from_blob function.\nvoid deleter(void* ptr)\n{\n printf(\"deleter(ptr=%p)\\n\",ptr);\n cudaFree(ptr);\n}\n*/\n\ntemplate \nat::Tensor blob_view(T* raw_ptr, std::vector shape, const at::TensorOptions& options, bool channels_last) {\n size_t size = 1;\n std::vector strides(shape.size());\n if (channels_last) {\n assert(shape.size() == 4);\n strides[0] = shape[1] * shape[2] * shape[3];\n strides[1] = 1;\n strides[2] = shape[1] * shape[3];\n strides[3] = shape[1];\n } else {\n int idx = strides.size();\n for (auto it = shape.rbegin(); it != shape.rend(); ++it) {\n strides[--idx] = size;\n size *= *it;\n }\n }\n size *= sizeof(T);\n // TODO: Implement dynamic reuse of pooled peer memory.\n // We provide no deleter function because all peer memory allocations are static in this implementation.\n return torch::from_blob((void*)raw_ptr, shape, strides, 0L, options);\n}\n\nvoid tensor_shape(at::Tensor t, bool explicit_nhwc, int& N, int& C, int& H, int& W) {\n if (t.dim() == 3) {\n N = 1;\n if (explicit_nhwc) {\n C = t.size(2);\n H = t.size(0);\n W = t.size(1);\n } else {\n C = t.size(0);\n H = t.size(1);\n W = t.size(2);\n }\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n N = t.size(0);\n C = t.size(3);\n H = t.size(1);\n W = t.size(2);\n } else {\n N = t.size(0);\n C = t.size(1);\n H = t.size(2);\n W = t.size(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, int(t.dim()));\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\nvoid tensor_strides(at::Tensor t, bool explicit_nhwc, int& stride_N, int& stride_C, int& stride_H, int& stride_W) {\n if (t.dim() == 3) {\n if (explicit_nhwc) {\n stride_C = t.stride(2);\n stride_H = t.stride(0);\n stride_W = t.stride(1);\n } else {\n stride_C = t.stride(0);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n }\n stride_N = t.size(0) * t.size(1) * t.size(2);\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n stride_N = t.stride(0);\n stride_C = t.stride(3);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n } else {\n stride_N = t.stride(0);\n stride_C = t.stride(1);\n stride_H = t.stride(2);\n stride_W = t.stride(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, t.dim());\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\ntemplate \ninline __device__ void __zero(T* dst) {\n *dst = T(0);\n}\n\ninline __device__ void __zero(int2* dst) { *dst = {0, 0}; }\n\ntemplate \ninline __device__ void zero_tensor(const int dim0, const int dim1, const int dim2, T* __restrict__ data,\n const int data_stride0, const int data_stride1, const int data_stride2,\n const int thread_id, const int block_id, const int num_blocks) {\n const int global_id = thread_id + block_id * THREADS_PER_CTA;\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i = global_id; i < count; i += num_threads) {\n int offset;\n if (contiguous) {\n offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n offset = j0 * data_stride0 + j1 * data_stride1 + j2 * data_stride2;\n }\n __zero(data + offset);\n }\n}\n\ntemplate \ninline __device__ void push_pull_tensor(const int dim0, const int dim1, const int dim2, const T* __restrict__ data_in,\n const int data_in_stride0, const int data_in_stride1, const int data_in_stride2,\n T* __restrict__ data_out, const int data_out_stride0,\n const int data_out_stride1, const int data_out_stride2, int4* local_peer,\n int4* remote_peer, const int thread_id, const int block_id,\n const int num_blocks) {\n // 128b=16B NVLink flit\n // Note: Use last 4B as a semaphore\n static_assert(sizeof(T) <= 12);\n union Flit {\n T payload;\n uint uints[4];\n };\n // Communication bit indicates whether flit has been received from\n // a remote GPU\n constexpr uint communication_mask = 1 << 0;\n // Status bit is used to choose the active peer buffer in an\n // alternating double buffer scheme. We use buffer 1 if the bits\n // match, use buffer 2 if the bits differ, and invert the bit\n // after finishing with a buffer.\n constexpr uint status_mask = 1 << 1;\n\n // Split peer memory into two sets of buffers\n // Note: Each block owns a THREADS_PER_CTA*2*16B chunk of peer\n // memory\n const int peer_offset1 = block_id * THREADS_PER_CTA * 2 + thread_id;\n const int peer_offset2 = peer_offset1 + THREADS_PER_CTA;\n volatile int* local_peer1 = reinterpret_cast(local_peer + peer_offset1);\n volatile int* local_peer2 = reinterpret_cast(local_peer + peer_offset2);\n volatile int* remote_peer1 = reinterpret_cast(remote_peer + peer_offset1);\n volatile int* remote_peer2 = reinterpret_cast(remote_peer + peer_offset2);\n\n // Iterate through tensor entries\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i0 = block_id * THREADS_PER_CTA; i0 < count; i0 += num_threads) {\n const int i = i0 + thread_id;\n const bool has_data = i < count;\n\n // Calculate buffer positions\n int data_in_offset, data_out_offset;\n if (contiguous) {\n data_in_offset = i;\n data_out_offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n data_in_offset = j0 * data_in_stride0 + j1 * data_in_stride1 + j2 * data_in_stride2;\n data_out_offset = j0 * data_out_stride0 + j1 * data_out_stride1 + j2 * data_out_stride2;\n }\n\n // Determine which peer memory buffer to use\n // Note: The status bit is not affected by asynchronous\n // communication from the remote GPU.\n Flit local_message1, local_message2;\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message1.uints[0]), \"=r\"(local_message1.uints[1]), \"=r\"(local_message1.uints[2]),\n \"=r\"(local_message1.uints[3])\n : \"l\"(local_peer1)\n : \"memory\");\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message2.uints[0]), \"=r\"(local_message2.uints[1]), \"=r\"(local_message2.uints[2]),\n \"=r\"(local_message2.uints[3])\n : \"l\"(local_peer2)\n : \"memory\");\n const uint status1 = local_message1.uints[3] & status_mask;\n const uint status2 = local_message2.uints[3] & status_mask;\n const bool peer1_is_active = (status1 ^ status2) == 0;\n volatile int* ox = peer1_is_active ? remote_peer1 : remote_peer2;\n volatile int* ix = peer1_is_active ? local_peer1 : local_peer2;\n const uint status = peer1_is_active ? status1 : status2;\n Flit recv_message = peer1_is_active ? local_message1 : local_message2;\n\n // Send flit to remote GPU\n // Note: Set communication bit and keep status bit\n Flit send_message;\n if (has_data) {\n send_message.payload = data_in[data_in_offset];\n }\n send_message.uints[3] = communication_mask | status;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ox), \"r\"(send_message.uints[0]),\n \"r\"(send_message.uints[1]), \"r\"(send_message.uints[2]), \"r\"(send_message.uints[3])\n : \"memory\");\n\n // Recieve flit from peer\n while ((recv_message.uints[3] & communication_mask) == 0) {\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(recv_message.uints[0]), \"=r\"(recv_message.uints[1]), \"=r\"(recv_message.uints[2]),\n \"=r\"(recv_message.uints[3])\n : \"l\"(ix)\n : \"memory\");\n }\n if (has_data) {\n data_out[data_out_offset] = recv_message.payload;\n }\n\n // Reset semaphore\n // Note: Clear communication bit and invert status bit\n uint flag = ~status & status_mask;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ix), \"n\"(0), \"n\"(0), \"n\"(0), \"r\"(flag)\n : \"memory\");\n if (i0 + num_threads < count) {\n __threadfence_system();\n }\n }\n}\n\ntemplate \n#if __CUDA_ARCH__ >= 700\n__launch_bounds__(THREADS_PER_CTA)\n#endif\n __global__ void push_pull_halos_1d_kernel(\n // top halo,\n T* toh, int toh_stride0, int toh_stride1, int toh_stride2, // top output halo (local)\n const T* tih, int tih_stride0, int tih_stride1, int tih_stride2, // top input halo (local)\n int4* tox, // top output transfer buffer (remote peer)\n int4* tix, // top input transfer buffer (local peer)\n // btm halo\n T* boh, int boh_stride0, int boh_stride1, int boh_stride2, // btm output halo (local)\n const T* bih, int bih_stride0, int bih_stride1, int bih_stride2, // btm input halo (local)\n int4* box, // btm output transfer buffer (remote peer)\n int4* bix, // btm input transfer buffer (local peer)\n // dimensions\n int dim0, int dim1, int dim2,\n bool top_first // whether to launch communicate top halo first\n ) {\n const int num_blocks_side = gridDim.x / 2;\n const int block_id_side = (blockIdx.x < num_blocks_side ? blockIdx.x : blockIdx.x - num_blocks_side);\n const bool in_top_block = top_first == (blockIdx.x < num_blocks_side);\n if (in_top_block) {\n if (top_zero) {\n zero_tensor(dim0, dim1, dim2, toh, toh_stride0, toh_stride1, toh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, tih, tih_stride0, tih_stride1, tih_stride2, toh, toh_stride0,\n toh_stride1, toh_stride2, tix, tox, threadIdx.x, block_id_side, num_blocks_side);\n }\n } else {\n if (btm_zero) {\n zero_tensor(dim0, dim1, dim2, boh, boh_stride0, boh_stride1, boh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, bih, bih_stride0, bih_stride1, bih_stride2, boh, boh_stride0,\n boh_stride1, boh_stride2, bix, box, threadIdx.x, block_id_side, num_blocks_side);\n }\n }\n}\n\n__global__ void delay_kernel(int delay_nanoseconds, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay_nanoseconds);\n *counter = new_counter;\n }\n}\n\n} // namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace peer_memory {\n\nint64_t allocate_raw(int64_t size) {\n float* ptr = 0L;\n cudaMalloc(&ptr, size);\n cudaMemset(ptr, 0, size);\n return (int64_t)ptr;\n}\n\nvoid free_raw(int64_t raw) { cudaFree((void*)raw); }\n\nvoid zero(int64_t raw, int64_t size) { cudaMemset((void*)raw, 0, size); }\n\nat::Tensor get_raw_ipc_address(int64_t raw) {\n cudaIpcMemHandle_t mem_handle;\n CUDACHECK(cudaIpcGetMemHandle(&mem_handle, (void*)raw));\n const int n = sizeof(cudaIpcMemHandle_t);\n auto address_tensor = torch::empty({n}, torch::dtype(torch::kUInt8));\n auto address_tensor_p = address_tensor.data_ptr();\n memcpy(address_tensor_p, (uint8_t*)&mem_handle, n);\n return address_tensor;\n}\n\nstd::vector get_raw_peers(at::Tensor ipc_addresses, int peer_rank, int64_t raw) {\n int peer_group_size = ipc_addresses.size(0);\n std::vector results(peer_group_size);\n for (int i = 0; i < peer_group_size; ++i) {\n if (i != peer_rank) {\n cudaIpcMemHandle_t mem_handle;\n memcpy(&mem_handle, ipc_addresses.index({i}).data_ptr(), sizeof(cudaIpcMemHandle_t));\n void* p = 0L;\n CUDACHECK(cudaIpcOpenMemHandle((void**)&p, mem_handle, cudaIpcMemLazyEnablePeerAccess));\n results[i] = (int64_t)p;\n } else {\n results[i] = (int64_t)raw;\n }\n }\n return results;\n}\n\nat::Tensor blob_view_half(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((at::Half*)raw, shape, torch::dtype(torch::kFloat16).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_float(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((float*)raw, shape, torch::dtype(torch::kFloat32).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_int(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((int*)raw, shape, torch::dtype(torch::kInt32).device(torch::kCUDA), channels_last);\n}\n\nvoid push_pull_halos_1d(\n bool diagnostics, bool explicit_nhwc,\n int numSM, // number of SMs to use (zero corresponds to all SMs)\n int rank, // rank in spatial parallel group\n bool top_zero, // if top halo should be zeroed\n at::Tensor top_in_halo, // top input halo buffer (in local device memory, sent to top neighbor)\n at::Tensor top_in_transfer, // top input transfer buffer (in local peer memory)\n at::Tensor top_out_transfer, // top output transfer buffer (in top neighbor peer memory)\n at::Tensor top_out_halo, // top output halo buffer (in local device memory, received from top neighbor)\n bool btm_zero, // if btm halo should be zeroed\n at::Tensor btm_in_halo, // btm input halo buffer (in local device memory, sent to btm neighbor)\n at::Tensor btm_in_transfer, // btm input transfer buffer (in local peer memory)\n at::Tensor btm_out_transfer, // btm output transfer buffer (in btm neighbor peer memory)\n at::Tensor btm_out_halo // btm output halo buffer (in local device memory, received from btm neighbor)\n) {\n // basic checks of inputs\n TORCH_CHECK(!(top_zero && btm_zero));\n TORCH_CHECK(top_in_halo.is_cuda());\n TORCH_CHECK(top_out_transfer.is_cuda());\n TORCH_CHECK(top_in_transfer.is_cuda());\n TORCH_CHECK(top_out_halo.is_cuda());\n TORCH_CHECK(btm_in_halo.is_cuda());\n TORCH_CHECK(btm_out_transfer.is_cuda());\n TORCH_CHECK(btm_in_transfer.is_cuda());\n TORCH_CHECK(btm_out_halo.is_cuda());\n\n // tensor shapes\n int tih_N, tih_C, tih_H, tih_W;\n tensor_shape(top_in_halo, explicit_nhwc, tih_N, tih_C, tih_H, tih_W);\n int toh_N, toh_C, toh_H, toh_W;\n tensor_shape(top_out_halo, explicit_nhwc, toh_N, toh_C, toh_H, toh_W);\n int bih_N, bih_C, bih_H, bih_W;\n tensor_shape(btm_in_halo, explicit_nhwc, bih_N, bih_C, bih_H, bih_W);\n int boh_N, boh_C, boh_H, boh_W;\n tensor_shape(btm_out_halo, explicit_nhwc, boh_N, boh_C, boh_H, boh_W);\n TORCH_CHECK(toh_N == tih_N && tih_N == boh_N && boh_N == bih_N && toh_C == tih_C && tih_C == boh_C &&\n boh_C == bih_C && toh_H == tih_H && tih_H == boh_H && boh_H == bih_H && toh_W == tih_W &&\n tih_W == boh_W && boh_W == bih_W);\n int NN = toh_N, NC = toh_C, NH = toh_H, NW = toh_W;\n if (diagnostics) {\n printf(\"rank %d: NN=%d, NC=%d, NH=%d, NW=%d\\n\", rank, NN, NC, NH, NW);\n }\n TORCH_CHECK(NN == 1);\n\n // tensor strides\n int tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W;\n tensor_strides(top_in_halo, explicit_nhwc, tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W);\n int toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W;\n tensor_strides(top_out_halo, explicit_nhwc, toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W);\n int bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W;\n tensor_strides(btm_in_halo, explicit_nhwc, bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W);\n int boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W;\n tensor_strides(btm_out_halo, explicit_nhwc, boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W);\n if (diagnostics) {\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n }\n\n // determine if nhwc\n bool is_nhwc = (toh_stride_C == 1);\n if (diagnostics) {\n printf(\"rank %d: is_nhwc = %s\\n\", rank, is_nhwc ? \"true\" : \"false\");\n }\n\n // determine if contiguous\n bool contiguous = true;\n if ((NN - 1) * toh_stride_N + (NC - 1) * toh_stride_C + (NH - 1) * toh_stride_H + (NW - 1) * toh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if ((NN - 1) * boh_stride_N + (NC - 1) * boh_stride_C + (NH - 1) * boh_stride_H + (NW - 1) * boh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if (!top_zero) {\n if (toh_stride_N != tih_stride_N || toh_stride_C != tih_stride_C || toh_stride_H != tih_stride_H ||\n toh_stride_W != tih_stride_W) {\n contiguous = false;\n }\n }\n if (!btm_zero) {\n if (boh_stride_N != bih_stride_N || boh_stride_C != bih_stride_C || boh_stride_H != bih_stride_H ||\n boh_stride_W != bih_stride_W) {\n contiguous = false;\n }\n }\n if (diagnostics) {\n printf(\"rank %d: contiguous = %s\\n\", rank, contiguous ? \"true\" : \"false\");\n }\n\n // determine whether to communicate top halo first\n bool top_first = rank % 2 != 0;\n if (diagnostics) {\n printf(\"rank %d: top_first = %s\\n\", rank, top_first ? \"true\" : \"false\");\n }\n\n // peer memory buffers\n int tox_size = top_out_transfer.numel() * top_out_transfer.element_size();\n int tix_size = top_in_transfer.numel() * top_in_transfer.element_size();\n int box_size = btm_out_transfer.numel() * btm_out_transfer.element_size();\n int bix_size = btm_in_transfer.numel() * btm_in_transfer.element_size();\n if (!top_zero) {\n TORCH_CHECK(top_out_transfer.is_contiguous());\n TORCH_CHECK(top_in_transfer.is_contiguous());\n TORCH_CHECK(tox_size == tix_size);\n }\n if (!btm_zero) {\n TORCH_CHECK(btm_out_transfer.is_contiguous());\n TORCH_CHECK(btm_in_transfer.is_contiguous());\n TORCH_CHECK(box_size == bix_size);\n }\n\n // figure out launch parameters\n int device;\n cudaGetDevice(&device);\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, device);\n if (numSM <= 0 || numSM > prop.multiProcessorCount) {\n numSM = prop.multiProcessorCount;\n }\n auto current_stream = at::cuda::getCurrentCUDAStream();\n dim3 block(THREADS_PER_CTA, 1, 1);\n\n // helper macros to launch templated kernel\n#define LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, TOP_ZERO, BTM_ZERO, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n /* kernel configuration */ \\\n int numBlocksPerSm; \\\n cudaOccupancyMaxActiveBlocksPerMultiprocessor( \\\n &numBlocksPerSm, push_pull_halos_1d_kernel, THREADS_PER_CTA, 0); \\\n dim3 grid(numSM * numBlocksPerSm, 1, 1); \\\n if (grid.x % 2 != 0) { \\\n /* require even number of blocks (half for top, half for bottom) */ \\\n grid.x -= 1; \\\n } \\\n if ((grid.x / 2) * THREADS_PER_CTA > NUM_ELEMENTS) { \\\n /* only need enough blocks to cover top and bottom halo elements */ \\\n grid.x = 2 * ((NUM_ELEMENTS + THREADS_PER_CTA - 1) / THREADS_PER_CTA); \\\n } \\\n if (!TOP_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > tox_size) { \\\n grid.x = 2 * (tox_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n if (!BTM_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > box_size) { \\\n grid.x = 2 * (box_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n TORCH_CHECK(grid.x >= 2); \\\n \\\n /* launch kernel */ \\\n cudaLaunchCooperativeKernel((void*)push_pull_halos_1d_kernel, grid, block, \\\n KERNEL_ARGS, 0, current_stream); \\\n } while (false)\n#define LAUNCH_PUSH_PULL_HALO_KERNEL(T, CONTIGUOUS, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n if (top_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, true, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else if (btm_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, true, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } \\\n } while (false)\n\n AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, top_out_halo.scalar_type(), \"push_pull_halos_1d_kernel\", [&] {\n if (diagnostics) {\n printf(\"rank %d: size(scalar_t) = %ld\\n\", rank, sizeof(scalar_t));\n }\n scalar_t* toh_p = top_out_halo.data_ptr();\n scalar_t* tih_p = top_in_halo.data_ptr();\n int4* tox_p = reinterpret_cast(top_out_transfer.data_ptr());\n int4* tix_p = reinterpret_cast(top_in_transfer.data_ptr());\n scalar_t* boh_p = btm_out_halo.data_ptr();\n scalar_t* bih_p = btm_in_halo.data_ptr();\n int4* box_p = reinterpret_cast(btm_out_transfer.data_ptr());\n int4* bix_p = reinterpret_cast(btm_in_transfer.data_ptr());\n if (diagnostics) printf(\"rank %d: choosing halo exchange kernel\\n\", rank);\n\n // do int2 vector loads if channel count permits\n if (contiguous && (NN * NH * NW * NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do contiguous int2 transfers\n if (diagnostics) {\n }\n toh_stride_N = toh_stride_H = toh_stride_W = toh_stride_C = 1;\n tih_stride_N = tih_stride_H = tih_stride_W = tih_stride_C = 1;\n boh_stride_N = boh_stride_H = boh_stride_W = boh_stride_C = 1;\n bih_stride_N = bih_stride_H = bih_stride_W = bih_stride_C = 1;\n NC = (NN * NH * NW * NC * sizeof(scalar_t)) / sizeof(int2);\n NN = NH = NW = 1;\n if (diagnostics) {\n printf(\"rank %d: launching contiguous int2 halo exchange kernel\\n\", rank);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NN * NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, true, kernel_args, num_elem);\n } else if (is_nhwc && (NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do strided int2 transfers\n int divisor = sizeof(int2) / sizeof(scalar_t);\n if (diagnostics) {\n printf(\"rank %d: launching strided int2 halo exchange kernel\\n\", rank);\n }\n toh_stride_N /= divisor;\n toh_stride_H /= divisor;\n toh_stride_W /= divisor;\n tih_stride_N /= divisor;\n tih_stride_H /= divisor;\n tih_stride_W /= divisor;\n boh_stride_N /= divisor;\n boh_stride_H /= divisor;\n boh_stride_W /= divisor;\n bih_stride_N /= divisor;\n bih_stride_H /= divisor;\n bih_stride_W /= divisor;\n NC /= divisor;\n if (diagnostics) {\n printf(\"rank %d: divisor=%d\\n\", rank, divisor);\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, false, kernel_args, num_elem);\n } else {\n // cannot do int2 transfers\n if (diagnostics) {\n printf(\"rank %d: launching non-int2 halo exchange kernel\\n\", rank);\n }\n int num_elem = NC * NH * NW;\n if (is_nhwc) {\n void* kernel_args[] = {&toh_p, &toh_stride_H, &toh_stride_W, &toh_stride_C, &tih_p, &tih_stride_H,\n &tih_stride_W, &tih_stride_C, &tox_p, &tix_p, &boh_p, &boh_stride_H,\n &boh_stride_W, &boh_stride_C, &bih_p, &bih_stride_H, &bih_stride_W, &bih_stride_C,\n &box_p, &bix_p, &NH, &NW, &NC, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n } else {\n void* kernel_args[] = {&toh_p, &toh_stride_C, &toh_stride_H, &toh_stride_W, &tih_p, &tih_stride_C,\n &tih_stride_H, &tih_stride_W, &tox_p, &tix_p, &boh_p, &boh_stride_C,\n &boh_stride_H, &boh_stride_W, &bih_p, &bih_stride_C, &bih_stride_H, &bih_stride_W,\n &box_p, &bix_p, &NC, &NH, &NW, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n }\n }\n });\n\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL_BASE\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL\n}\n\n} // namespace peer_memory\n} // namespace contrib\n} // namespace apex\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/peer_memory/peer_memory_cuda.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \n#include \n\n#ifdef OLD_GENERATOR_PATH\n#include \n#else\n#include \n#endif\n\n#include \n#include \n\n#include \n\n#include \"philox.cuh\"\n\n// Warp reduce kernels to reduce N groups of data into N numbers, where N = warpSize / width.\n// width should be a power of 2 and should be less than warpSize.\ntemplate \n__device__ __forceinline__ scalar_t warpReduce(scalar_t x, int width = C10_WARP_SIZE) {\n for (unsigned offset = width / 2; offset > 0; offset /= 2) {\n x += __shfl_down_sync(0xffffffff, x, offset, width);\n }\n return x;\n}\n\ninline int largestPowerOfTwo(int x) {\n int y = 1;\n while (y <= x) y <<= 1;\n return y >> 1;\n}\n\n/*\nFigure out vectorization type for masks.\nSimilar to how PyTorch figures out acc_t here:\naten/src/ATen/AccumulateType.h\n*/\ntemplate \nstruct MaskVecType {};\n\ntemplate <>\nstruct MaskVecType<1> {\n using type = uint8_t;\n};\ntemplate <>\nstruct MaskVecType<2> {\n using type = uint16_t;\n};\ntemplate <>\nstruct MaskVecType<4> {\n using type = uint32_t;\n};\n\ntemplate \nusing mvec_type = typename MaskVecType::type;\n\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels.\n// For fwd, batch offset and stride are different for packing and non-packing mode.\nstruct OffsetCalFwd {\n __device__ __forceinline__ OffsetCalFwd(int64_t batch, const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t gLen, int64_t hiddenSize, bool packOutput)\n : batch(batch),\n batchOffset(batchOffset),\n maxFLen(maxFLen),\n maxGLen(maxGLen),\n gLen(gLen),\n hiddenSize(hiddenSize),\n packOutput(packOutput) {}\n\n int64_t batch;\n const int64_t* batchOffset;\n int64_t maxFLen;\n int64_t maxGLen;\n int64_t gLen;\n int64_t hiddenSize;\n bool packOutput;\n\n __device__ __forceinline__ int64_t getBatchOffset() {\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\n : batch * maxFLen * maxGLen * hiddenSize;\n }\n\n __device__ __forceinline__ int64_t getStrideF() { return packOutput ? gLen * hiddenSize : maxGLen * hiddenSize; }\n};\n\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels\n// For bwd, batch offset and stride are different for packing and non-packing mode.\n// The reducion is done for two input tensors. Therefore, generating two sets of offsets\n// according to bwdFasterDim can lead to a unified implementation in the actual kernel.\nstruct OffsetCalBwd {\n __device__ __forceinline__ OffsetCalBwd(int64_t batch, const int64_t* batchOffset, const int* fLen, const int* gLen,\n int64_t maxFLen, int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim)\n : batch(batch),\n batchOffset(batchOffset),\n maxFLen(maxFLen),\n maxGLen(maxGLen),\n fLen(fLen),\n gLen(gLen),\n hiddenSize(hiddenSize),\n packOutput(packOutput),\n bwdFasterDim(bwdFasterDim) {}\n\n int64_t batch;\n const int64_t* batchOffset;\n const int* fLen;\n const int* gLen;\n int64_t maxFLen;\n int64_t maxGLen;\n int64_t hiddenSize;\n bool packOutput;\n bool bwdFasterDim; // whether doing bwd on the faster moving dimension\n\n __device__ __forceinline__ int64_t getBatchOffset() {\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\n : batch * maxFLen * maxGLen * hiddenSize;\n }\n\n __device__ __forceinline__ int64_t getMaxXLen() { return bwdFasterDim ? maxGLen : maxFLen; }\n\n __device__ __forceinline__ auto getMyXLen() -> decltype(gLen[batch]) {\n return bwdFasterDim ? gLen[batch] : fLen[batch];\n }\n\n __device__ __forceinline__ auto getMyYLen() -> decltype(gLen[batch]) {\n return bwdFasterDim ? fLen[batch] : gLen[batch];\n }\n\n __device__ __forceinline__ int64_t getStrideX() {\n return bwdFasterDim ? hiddenSize : ((packOutput ? gLen[batch] : maxGLen) * hiddenSize);\n }\n\n __device__ __forceinline__ int64_t getStrideY() {\n return bwdFasterDim ? ((packOutput ? gLen[batch] : maxGLen) * hiddenSize) : hiddenSize;\n }\n};\n\n// Vanila transducer joint forward kernel\n// Detail of this joint function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n\n// f is a tensor of shape [batch, T, H]\n// g is a tensor of shape [batch, U, H]\n// the transducer joint does\n// sum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\n// The resultant tensor is of shape [batch, T, U, H]\n// Each thread block is working on one \"batch\" of data in the output tensor, [batch, t, u, :]\n\n// This joint function can optionally pack the output where the output tensor with a shape of\n// [B, T, U, H] is packed into [B_packed, H].\n// Don't-care region (t > fLen) or (u > gLen) is removed.\n// To enable packing, the starting offset for each batch need to be specified with batchOffset.\ntemplate \n__global__ void transducer_joint_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t hiddenSize, bool packOutput, scalar_t* sum) {\n const int batch = blockIdx.z;\n const int t = blockIdx.y;\n const int u = blockIdx.x;\n const auto myFLen = fLen[batch];\n const auto myGLen = gLen[batch];\n\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideF = offsetCal.getStrideF();\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize;\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize;\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize;\n\n if (t < myFLen and u < myGLen) {\n#pragma unroll\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\n if (h < hiddenSize) {\n mySum[h] = myF[h] + myG[h];\n }\n }\n } else if (packOutput == false and t < maxFLen and u < maxGLen) {\n// Need to write finite data to don't-care region because we instantiate the result tensor\n// with torch::empty for performance reasons. Even though it is don't-care region, the\n// contents need to be finite, otherwise could lead to NaN in WGRAD.\n// In packing mode, this write is no longer necessary as we remove the don't-care region\n// from the output.\n// Picking -1 (over 0) here for ease of testing.\n#pragma unroll\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\n if (h < hiddenSize) {\n mySum[h] = -1;\n }\n }\n }\n}\n\n/*\nTiled version of the joint forward kernel\nDetail of this joint function can be found in:\n[1] Sequence Transduction with Recurrent Neural Networks.\n\nf is a tensor of shape [batch, T, H]\ng is a tensor of shape [batch, U, H]\nthe transducer joint does\nsum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\nThe resultant tensor is of shape [batch, T, U, H]\nEach thread is working on a tile of the shape of tileF x tileG in the result tensor.\nThe input for the tile is first loaded in the register and is reused tileG and tileF times.\n\nThis joint function can optionally pack the output where the output tensor with a shape of\n[B, T, U, H] is packed into [B_packed, H].\nDon't-care region (t > fLen) or (u > gLen) is removed.\nTo enable packing, the starting offset for each batch need to be specified with batchOffset.\n\nOptionally this joint function performs ReLU and/or dropout on the joint output, which is\ncontrolled by arguments relu and dropout, respectively. philoxArgs is argument used for generating\npseudorandom number. When at least one of operations in ReLU and dropout is activated, the joint\nfunction is a masked operation, which is controlled by the template argument masked. In this case,\nmasks are saved to backward.\n*/\ntemplate \n__global__ void transducer_joint_tiled_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\n int64_t hiddenSize, int64_t hiddenPerBlock, bool packOutput, bool relu,\n bool dropout, float p, at::PhiloxCudaState philoxArgs, scalar_t* sum,\n uint8_t* mask) {\n static_assert(U == 4, \"U has to be 4, as random numbers are generated in batch of 4\");\n\n const int batch = blockIdx.z;\n const int t = blockIdx.y * tileF;\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\n const int u = blockIdx.x / hiddenBlock * tileG;\n const int hOffset = (blockIdx.x % hiddenBlock) * hiddenPerBlock;\n const int h = threadIdx.x;\n const auto myFLen = fLen[batch];\n const auto myGLen = gLen[batch];\n\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideF = offsetCal.getStrideF();\n\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize + hOffset;\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize + hOffset;\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\n uint8_t* myMask = mask + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\n\n // The following code is only needed for dropout. We try to bypass them as much as possible.\n auto seeds = masked ? at::cuda::philox::unpack(philoxArgs)\n : std::make_tuple(static_cast(0), static_cast(0));\n uint64_t tid =\n masked ? (static_cast(blockIdx.z) * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + blockIdx.x) *\n blockDim.x +\n threadIdx.x\n : 0;\n Philox ph(std::get<0>(seeds), tid, std::get<1>(seeds));\n scalar_t scale = masked ? ((p == 0) ? 0 : 1 / p) : 0;\n bool dropoutMask[U];\n\n if (t < myFLen and u < myGLen and hOffset + h < hiddenSize) {\n // register buffers for tiled input reuse\n scalar_t fBuffer[tileF], gBuffer[tileG];\n for (int i = 0; i < tileF; ++i) {\n if (t + i < myFLen) fBuffer[i] = myF[i * hiddenSize + h];\n }\n for (int j = 0; j < tileG; ++j) {\n if (u + j < myGLen) gBuffer[j] = myG[j * hiddenSize + h];\n }\n#pragma unroll\n for (int i = 0; i < tileF; ++i) {\n if (t + i < myFLen) {\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n int idx = i * tileG + j;\n if (masked and dropout and idx % U == 0) {\n // For performance, generate 4 random numbers in one shot\n // auto rand4 = curand_uniform4(&state);\n auto rand4 = uniform4(ph());\n dropoutMask[0] = rand4.x < p;\n dropoutMask[1] = rand4.y < p;\n dropoutMask[2] = rand4.z < p;\n dropoutMask[3] = rand4.w < p;\n }\n\n if (u + j < myGLen) {\n scalar_t out = fBuffer[i] + gBuffer[j];\n if (masked) {\n // Apply ReLU here when relu is True\n bool localMask = relu ? (out > 0) : 1;\n localMask = dropout ? localMask & dropoutMask[idx % U] : localMask;\n out = dropout ? out * localMask * scale : out * localMask;\n myMask[i * strideF + j * hiddenSize + h] = static_cast(localMask);\n }\n mySum[i * strideF + j * hiddenSize + h] = out;\n } else if (packOutput == false and u + j < maxGLen)\n mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n } else if (packOutput == false and t + i < maxFLen) {\n// Again need to write finite data to don't-care region\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n }\n }\n } else if (packOutput == false and t < maxFLen and u < maxGLen and hOffset + h < hiddenSize) {\n// Only need to ensure the finity in normal mode\n#pragma unroll\n for (int i = 0; i < tileF; ++i) {\n if (t + i < maxFLen) {\n#pragma unroll\n for (int j = 0; j < tileG; ++j) {\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\n }\n }\n }\n }\n}\n\n/*\nBwd operation (reduction) on one input tensor. Since the operation performed for the two input\ntensors are exactly the same, only one kernel is needed, and the different indexing offsets\nand strides are handled by OffsetCalBwd.\n\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\nnon-packed form.\n\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__device__ void transducer_joint_single_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim, // whether bwd on the faster moving dimension (u)\n float scale, scalar_t* inGrad, int yBlockOffset = 0) {\n const int batch = blockIdx.z;\n // For the second input tensor, this offset need to be subtracted because the first yBlockOffset\n // sets of thread blocks are for the first input tensor.\n const int x = blockIdx.y - yBlockOffset;\n const int hOffset = blockIdx.x * C10_WARP_SIZE;\n const int wid = threadIdx.y;\n const int lid = threadIdx.x;\n const int numWarp = blockDim.y;\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\n const auto maxXLen = offsetCal.getMaxXLen();\n const auto myXLen = offsetCal.getMyXLen();\n const auto myYLen = offsetCal.getMyYLen();\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\n\n if (x < myXLen) {\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideX = offsetCal.getStrideX();\n const auto strideY = offsetCal.getStrideY();\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\n\n // Each warp reduces numYPerWarp \"y\" first\n acc_t warpSum = 0;\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\n#pragma unroll\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\n auto y = wid * numYPerWarp + warpY;\n if (y < myYLen and (hOffset + lid) < hiddenSize)\n if (masked)\n warpSum += static_cast(myGrad[y * strideY + lid]) * myMask[y * strideY + lid] * scale;\n else\n warpSum += myGrad[y * strideY + lid];\n }\n\n // transpose partial sum in SMEM and reduce further using warpReduce\n smem[lid * numWarp + wid] = warpSum;\n __syncthreads();\n auto sum = smem[wid * C10_WARP_SIZE + lid];\n sum = warpReduce(sum, numWarp);\n\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // example of 4 warps (a, b, c, d) with 8 threads per warp\n // Each warp need 8 / 4 = 2 threads to write the results.\n if (hOffset + wid * C10_WARP_SIZE / numWarp + lid / numWarp < hiddenSize) {\n if (lid % numWarp == 0) {\n myInGrad[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = sum;\n }\n }\n } else if (wid == 0 and hOffset + lid < hiddenSize) {\n // Need to ensure the grad is zero for don't care region\n myInGrad[lid] = 0;\n }\n}\n\n/*\nActual bwd (reduction) kernel get launched.\nCall transducer_joint_single_backward twice on two input tensors.\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\nuses the rest.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__global__ void transducer_joint_combined_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput, float scale,\n scalar_t* fGrad, scalar_t* gGrad) {\n if (blockIdx.y < maxFLen) {\n transducer_joint_single_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\n } else {\n transducer_joint_single_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\n }\n}\n\n/*\nVectorized version of transducer_joint_single_backward\nDoing exact same operation as transducer_joint_single_backward except the load and store are\nvectorized.\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\nnon-packed form.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__device__ void transducer_joint_single_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n bool bwdFasterDim, float scale, scalar_t* inGrad,\n int yBlockOffset = 0) {\n const int batch = blockIdx.z;\n const int x = blockIdx.y - yBlockOffset;\n const int hOffset = blockIdx.x * C10_WARP_SIZE * V;\n const int wid = threadIdx.y;\n const int lid = threadIdx.x;\n const int numWarp = blockDim.y;\n\n // Figure out the vectorization type for mask\n using mvec_t = mvec_type;\n\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\n const auto maxXLen = offsetCal.getMaxXLen();\n const auto myXLen = offsetCal.getMyXLen();\n const auto myYLen = offsetCal.getMyYLen();\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n acc_t warpSum[V];\n scalar_t inBuffer[V];\n uint8_t maskBuffer[V];\n scalar_t outBuffer[V];\n auto myInGradVec = reinterpret_cast(myInGrad);\n auto outBufferVec = reinterpret_cast(outBuffer);\n\n if (x < myXLen) {\n const auto myBatchOffset = offsetCal.getBatchOffset();\n const auto strideX = offsetCal.getStrideX();\n const auto strideY = offsetCal.getStrideY();\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\n\n for (int i = 0; i < V; ++i) warpSum[i] = 0;\n\n // Each warp reduces numYPerWarp \"y\" first\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\n auto y = wid * numYPerWarp + warpY;\n auto myGradVec = reinterpret_cast(myGrad + y * strideY);\n auto myMaskVec = masked ? reinterpret_cast(myMask + y * strideY) : nullptr;\n auto inBufferVec = reinterpret_cast(inBuffer);\n auto maskBufferVec = reinterpret_cast(maskBuffer);\n if (hOffset + lid * V < hiddenSize and y < myYLen) {\n *inBufferVec = myGradVec[lid]; // vectorized load\n if (masked) {\n *maskBufferVec = myMaskVec[lid];\n#pragma unroll\n for (int i = 0; i < V; ++i) warpSum[i] += static_cast(inBuffer[i]) * maskBuffer[i] * scale;\n } else {\n#pragma unroll\n for (int i = 0; i < V; ++i) warpSum[i] += inBuffer[i];\n }\n }\n }\n\n // transpose partial sum in SMEM and reduce further using warpReduce\n for (int i = 0; i < V; ++i) {\n smem[lid * numWarp + wid] = warpSum[i];\n __syncthreads();\n auto sum = smem[wid * C10_WARP_SIZE + lid];\n\n if (hOffset + (wid * C10_WARP_SIZE / numWarp) * V < hiddenSize) {\n sum = warpReduce(sum, numWarp);\n if (lid % numWarp == 0) {\n outBuffer[i] = sum;\n }\n }\n __syncthreads();\n }\n\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // a a b b c c d d\n // example of 4 warps (a, b, c, d) with 8 threads per warp\n // Each warp need 8 / 4 = 2 threads to write the results.\n if (lid % numWarp == 0 and hOffset + (wid * C10_WARP_SIZE / numWarp + lid / numWarp) * V < hiddenSize)\n myInGradVec[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = *outBufferVec;\n } else if (wid == 0 and hOffset + lid * V < hiddenSize) {\n // Need to ensure the grad is zero for don't care region\n myInGradVec[lid] = 0;\n }\n}\n\n/*\nVecotrized version of transducer_joint_combined_backward\nCall transducer_joint_single_vec_backward twice on two input tensors.\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\nuses the rest.\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\nand mask contains the mask information.\n*/\ntemplate \n__global__ void transducer_joint_combined_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\n float scale, scalar_t* fGrad, scalar_t* gGrad) {\n if (blockIdx.y < maxFLen) {\n transducer_joint_single_vec_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\n } else {\n transducer_joint_single_vec_backward(\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\n }\n}\n\nstd::vector transducer_joint_cuda_forward(torch::Tensor f, torch::Tensor g, torch::Tensor fLen,\n torch::Tensor gLen, torch::Tensor batchOffset,\n int64_t packedBatch, int opt, bool packOutput, bool relu,\n bool dropout, float dropoutProb, int tileSize) {\n auto tensorOpt = f.options();\n auto dtype = f.scalar_type();\n const auto batchSize = f.size(0);\n const auto maxFLen = f.size(1);\n const auto maxGLen = g.size(1);\n const auto hiddenSize = f.size(2);\n bool masked = dropout or relu;\n\n int64_t* batchOffsetPtr = nullptr;\n torch::Tensor sum, mask;\n auto maskOpt = tensorOpt.dtype(torch::kUInt8);\n if (!packOutput) {\n sum = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, tensorOpt);\n batchOffsetPtr = nullptr;\n if (masked) mask = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, maskOpt);\n } else {\n sum = torch::empty({packedBatch, hiddenSize}, tensorOpt);\n batchOffsetPtr = batchOffset.data_ptr();\n if (masked) mask = torch::empty({packedBatch, hiddenSize}, maskOpt);\n }\n uint8_t* maskPtr = masked ? mask.data_ptr() : nullptr;\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n TORCH_CHECK(opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\n // Simple heuristics\n const int numThread =\n std::min(128, (static_cast(hiddenSize) + C10_WARP_SIZE - 1) / C10_WARP_SIZE * C10_WARP_SIZE);\n\n if (opt == 0) {\n // vanilla kernel\n const int threads = numThread;\n const dim3 blocks(maxGLen, maxFLen, batchSize);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_forward\", ([&] {\n transducer_joint_forward<<>>(\n f.data_ptr(), g.data_ptr(), fLen.data_ptr(), gLen.data_ptr(),\n batchOffsetPtr, maxFLen, maxGLen, hiddenSize, packOutput, sum.data_ptr());\n }));\n }\n if (opt == 1) {\n // tiled version. For simplicity, assume tileF == tileG, even though the kernel can\n // support more general cases.\n const int threads = numThread;\n const int hiddenPerBlock = numThread;\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\n const dim3 blocks((maxGLen + tileSize - 1) / tileSize * hiddenBlock, (maxFLen + tileSize - 1) / tileSize,\n batchSize);\n\n TORCH_CHECK(tileSize == 1 or tileSize == 2 or tileSize == 4, \"Expected tileSize to be in [1, 2, 4], but got \",\n tileSize);\n\n at::PhiloxCudaState rng_engine_inputs;\n if (masked) {\n // set up PRG when the input is masked. rng_engine_inputs will be used as a space filler\n // for non-masked calls.\n // Therefore no need to initialize.\n c10::optional gen_;\n auto gen = at::get_generator_or_default(gen_, at::cuda::detail::getDefaultCUDAGenerator());\n // counterOffset records how many cuRAND calls each thread makes. For a tiled kernel,\n // each thread processes tileF * tileG output elements.\n int64_t counterOffset = tileSize * tileSize;\n {\n std::lock_guard lock(gen->mutex_);\n rng_engine_inputs = gen->philox_cuda_state(counterOffset);\n }\n }\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_forward\", ([&] {\n void (*kernel)(const scalar_t*, const scalar_t*, const int*, const int*, const int64_t*, int64_t, int64_t,\n int64_t, int64_t, bool, bool, bool, float, at::PhiloxCudaState, scalar_t*, uint8_t*);\n if (masked) {\n switch (tileSize) {\n case 2:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 4:\n kernel = &transducer_joint_tiled_forward;\n break;\n }\n } else {\n switch (tileSize) {\n case 1:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 2:\n kernel = &transducer_joint_tiled_forward;\n break;\n case 4:\n kernel = &transducer_joint_tiled_forward;\n break;\n }\n }\n\n kernel<<>>(f.data_ptr(), g.data_ptr(), fLen.data_ptr(),\n gLen.data_ptr(), batchOffsetPtr, maxFLen, maxGLen, hiddenSize,\n hiddenPerBlock, packOutput, relu, dropout, 1.0f - dropoutProb,\n rng_engine_inputs, sum.data_ptr(), maskPtr);\n }));\n }\n\n C10_CUDA_CHECK(cudaGetLastError());\n if (masked)\n return {sum, mask};\n else\n return {sum};\n}\n\nstd::vector transducer_joint_cuda_backward(std::vector in, torch::Tensor fLen,\n torch::Tensor gLen, torch::Tensor batchOffset, int maxFLen,\n int maxGLen, bool packOutput, float scale) {\n auto grad = in[0];\n bool masked = (in.size() == 2);\n uint8_t* maskPtr = masked ? in[1].data_ptr() : nullptr;\n\n auto tensorOpt = grad.options();\n auto dtype = grad.scalar_type();\n const int batchSize = fLen.size(0);\n const int hiddenSize = grad.size(-1);\n\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const int maxNumWarp = deviceProperties->maxThreadsPerBlock / C10_WARP_SIZE;\n\n torch::Tensor fGrad = torch::empty({batchSize, maxFLen, hiddenSize}, tensorOpt);\n torch::Tensor gGrad = torch::empty({batchSize, maxGLen, hiddenSize}, tensorOpt);\n\n int64_t* batchOffsetPtr = (!packOutput) ? nullptr : batchOffset.data_ptr();\n\n // The number \"y\" I would like each thread to work on\n const int workPerThread = 32;\n // Since the bwd for f and g have the same thread block size, we need to use the max of the two.\n int numWarp = largestPowerOfTwo((std::max(maxFLen, maxGLen) + workPerThread - 1) / workPerThread);\n // Would like to have at least 2 warps\n numWarp = std::max(2, numWarp);\n // cap on the maximum number of warps allowed\n numWarp = std::min(maxNumWarp, numWarp);\n\n // Need smem for transposing the partial sum. The partial sum is in a matrix of the shape\n // numWarp x warpSize\n const int smemSize = numWarp * C10_WARP_SIZE;\n const dim3 threads(C10_WARP_SIZE, numWarp, 1);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_joint_cuda_backward_kernel\", ([&] {\n auto gradPtr = grad.data_ptr();\n auto fLenPtr = fLen.data_ptr();\n auto gLenPtr = gLen.data_ptr();\n auto fGradPtr = fGrad.data_ptr();\n auto gGradPtr = gGrad.data_ptr();\n\n // resolve the acc_t type\n using acc_t = at::acc_type;\n using vec_t = uint64_t;\n\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\n constexpr int vecAlignment = std::alignment_of::value;\n\n // if all input and output tensors meet the alignment requirement\n bool memAlign = (reinterpret_cast(gradPtr) % vecAlignment == 0) and\n (reinterpret_cast(fGradPtr) % vecAlignment == 0) and\n (reinterpret_cast(gGradPtr) % vecAlignment == 0);\n\n if (vectFactor > 1 and hiddenSize % vectFactor == 0 and memAlign) {\n // If vectorization helps and the alignment requirement is met, use the vectorized\n // kernel. For simplicity, hiddenSize needs to be a multiple vecFactor.\n const dim3 blocks((hiddenSize + C10_WARP_SIZE * vectFactor - 1) / (C10_WARP_SIZE * vectFactor),\n maxFLen + maxGLen, batchSize);\n if (masked) {\n transducer_joint_combined_vec_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n } else {\n transducer_joint_combined_vec_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n }\n } else {\n const dim3 blocks((hiddenSize + C10_WARP_SIZE - 1) / C10_WARP_SIZE, maxFLen + maxGLen, batchSize);\n if (masked) {\n transducer_joint_combined_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n } else {\n transducer_joint_combined_backward\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\n fGradPtr, gGradPtr);\n }\n }\n }));\n\n return {fGrad, gGrad};\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/transducer/transducer_joint_kernel.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\ntemplate \n__device__ __forceinline__ scalar_t logSumExp(scalar_t a, scalar_t b) {\n // standard log-sum-exp trick is used here to provide better numerical stability\n return (a >= b) ? a + std::log1p(exp(b - a)) : b + std::log1p(exp(a - b));\n}\n\n// Vanilla transducer loss function (i.e. forward-backward algorithm)\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_forward(const scalar_t* x, const int* label, const int* audLen, const int* txtLen,\n const int64_t* batchOffset,\n int64_t dictSize, // 64-bit indexing for data tensor\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n const int tid = threadIdx.x;\n const auto myFLen = audLen[batch];\n // Note that start of the sentence is added as 1 here\n const auto myGLen = txtLen[batch] + 1;\n const auto myLabel = label + batch * (maxGLen - 1);\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n int u = tid;\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n if (u == 0) myAlpha[0] = 0;\n __syncthreads();\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; ++step) {\n // Move along the diagonal wavefront to leverage available parallelism\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0) {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u)\n myAlpha[t * maxGLen + u] = myAlpha[(t - 1) * maxGLen] + myX[((t - 1) * myStrideT) * dictSize + blankIdx];\n } else if (t == 0) {\n // alpha(t, u-1) = alpha(t, u-1) * y(t, u-1)\n myAlpha[u] = myAlpha[u - 1] + myX[(u - 1) * dictSize + myLabel[u - 1]];\n } else {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u) + alpha(t, u-1) * y(t, u-1)\n acc_t current = myAlpha[(t - 1) * maxGLen + u] + myX[((t - 1) * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myAlpha[t * maxGLen + u - 1] + myX[(t * myStrideT + u - 1) * dictSize + myLabel[u - 1]];\n myAlpha[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n if (u == 0) {\n myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n }\n __syncthreads();\n\n for (int64_t step = myFLen + myGLen - 3; step >= 0; --step) {\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1) {\n // beta(t, u) = beta(t+1, u) * null(t, u)\n myBeta[t * maxGLen + u] = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n } else if (t == myFLen - 1) {\n // beta(t, u) = beta(t, u+1) * y(t, u)\n myBeta[t * maxGLen + u] = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n } else {\n // beta(t, u) = beta(t+1, u)*null(t, u) + beta(t, u+1)*y(t, u)\n acc_t current = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n myBeta[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n if (tid == 0) loss[batch] = -myBeta[0];\n }\n}\n\n// transudcer loss function (i.e. forward-backward algorithm) with batch loading optimization.\n// Compared to the vanilla version, there are two optimizations:\n// 1. load x in batch through loop unrolling to reduce the latency.\n// 2. Use registers and shared memory to hold alpha and beta values passed from one step the next.\n// For simplicity, this kernel currently only supports U <= maxThread, which should be the common\n// case. For cases where U > maxThread, the vanilla kernel is used as a fallback option.\n\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_batch_load_forward(const scalar_t* x, const int* label, const int* audLen,\n const int* txtLen, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n int u = threadIdx.x;\n const auto myFLen = audLen[batch];\n const auto myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n scalar_t next[batchLdSize], current[batchLdSize];\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedAlpha[2] = {smem, smem + maxGLen};\n\n sharedAlpha[0][u] = 0;\n __syncthreads();\n\n if (u == 0) myAlpha[0] = 0;\n\n auto myAlphaLabel = (u == 0) ? 0 : label[batch * (maxGLen - 1) + u - 1];\n // register used to pass value to the next step for the same thread\n acc_t prvStepAlpha = 0;\n for (int64_t step = 1; step < myFLen + myGLen - 1 + batchLdSize; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X through loop unrolling\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = step + i - u;\n int64_t currentId = ((t - 1) * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u - 1) * dictSize + myAlphaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == 0) {\n current[i] = myX[currentId];\n } else if (t == 0) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedAlphaRd = sharedAlpha[(step + i - 1) % 2];\n auto sharedAlphaWr = sharedAlpha[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = step + i - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0)\n prvStepAlpha = prvStepAlpha + current[i];\n else if (t == 0)\n prvStepAlpha = sharedAlphaRd[u - 1] + next[i];\n else\n prvStepAlpha = logSumExp(prvStepAlpha + current[i], sharedAlphaRd[u - 1] + next[i]);\n sharedAlphaWr[u] = prvStepAlpha;\n myAlpha[t * maxGLen + u] = prvStepAlpha;\n }\n }\n __syncthreads();\n }\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedBeta[2] = {smem, smem + maxGLen};\n sharedBeta[0][u] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n __syncthreads();\n\n auto myBetaLabel = (u == maxGLen - 1) ? 0 : label[batch * (maxGLen - 1) + u];\n // register used to pass value to the next step for the same thread\n acc_t prvStepBeta = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n if (u == 0) myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = prvStepBeta;\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n int64_t currentId = (t * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u) * dictSize + myBetaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == myGLen - 1) {\n current[i] = myX[currentId];\n } else if (t == myFLen - 1) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedBetaRd = sharedBeta[(step + i - 1) % 2];\n auto sharedBetaWr = sharedBeta[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1)\n prvStepBeta = prvStepBeta + current[i];\n else if (t == myFLen - 1)\n prvStepBeta = sharedBetaRd[u + 1] + next[i];\n else\n prvStepBeta = logSumExp(prvStepBeta + current[i], sharedBetaRd[u + 1] + next[i]);\n sharedBetaWr[u] = prvStepBeta;\n myBeta[t * maxGLen + u] = prvStepBeta;\n }\n }\n __syncthreads();\n }\n }\n if (u == 0) loss[batch] = -prvStepBeta;\n }\n}\n\n// Vanilla transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// For this backward kernel, bwd op for the preceding softmax is assumed to be handled elsewhere,\n// hence only Eq(20) in [1] is implemented in this kernel.\n\n// Each thread block works on [batch, t, :, :] of data. Each thread works on a specific u at a time\n// Since only gradients for the correct token and null token need to be updated, gradients at other\n// locations are initialized to 0.\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha, const acc_t* beta,\n const int64_t* batchOffset, int64_t dictSize, int64_t blankIdx,\n int64_t maxFLen, int64_t maxGLen, bool packedInput, scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int t = blockIdx.x;\n const int batch = blockIdx.y;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n auto myX = x + (myBatchOffset + t * myStrideT) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT) * dictSize;\n auto myLabel = label + batch * (maxGLen - 1);\n\n int64_t u = tid;\n while (t < myFLen and u < myGLen) {\n // Do the update\n // loss = -ln(Pr(y*|x))\n acc_t grad = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n if (u != myGLen - 1)\n myXGrad[u * dictSize + myLabel[u]] =\n -std::exp(grad + myBeta[t * maxGLen + u + 1] + myX[u * dictSize + myLabel[u]]);\n if (t == myFLen - 1 and u == myGLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myX[u * dictSize + blankIdx]);\n else if (t != myFLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myBeta[(t + 1) * maxGLen + u] + myX[u * dictSize + blankIdx]);\n\n u += blockDim.x;\n }\n}\n\n// Fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n\n if (t < myFLen and u < myGLen) {\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n myLabelShared = myLabel[u];\n }\n\n __syncthreads();\n\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n // Do the update\n acc_t grad = commonFactor + myX[h]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGrad[h] = myGrad;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n myXGrad[h] = 0;\n }\n }\n}\n\n// Vectorized version of fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_vec_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // Variabels for vectorization\n scalar_t myXBuffer[V], myXGradBuffer[V];\n auto myXVec = reinterpret_cast(myX);\n auto myXGradVec = reinterpret_cast(myXGrad);\n auto myXBufferVec = reinterpret_cast(myXBuffer);\n auto myXGradBufferVec = reinterpret_cast(myXGradBuffer);\n if (t < myFLen and u < myGLen) {\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n if (t != myFLen - 1) myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n if (u != myGLen - 1) {\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myLabelShared = myLabel[u];\n }\n }\n\n __syncthreads();\n\n#pragma unroll\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n // Load myX in a vector form\n *myXBufferVec = myXVec[h0 / V];\n// Do the update for a vector of input\n#pragma unroll\n for (int i = 0; i < V; ++i) {\n auto h = h0 + i;\n acc_t grad = commonFactor + myXBuffer[i]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGradBuffer[i] = myGrad;\n }\n\n // Store myXGrad in a vector form\n myXGradVec[h0 / V] = *myXGradBufferVec;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n myXGradVec[h0 / V] = 0;\n }\n }\n}\n\nstd::vector transducer_loss_cuda_forward(torch::Tensor x, torch::Tensor label, torch::Tensor audLen,\n torch::Tensor txtLen, torch::Tensor batchOffset, int maxFLen,\n int blankIdx, int opt, bool packedInput) {\n auto scalarType = x.scalar_type();\n auto tensorOpt = x.options();\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n\n TORCH_CHECK(blankIdx >= 0 and blankIdx < dictSize, \"Expected blank index to be in the range of 0 to \", dictSize - 1,\n \", but got \", blankIdx);\n TORCH_CHECK(opt == -1 or opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\n\n // The data type of alpha and beta will be resolved at dispatch time,\n // hence defined here and assigned later\n torch::Tensor alpha;\n torch::Tensor beta;\n torch::Tensor loss = torch::empty({batchSize}, tensorOpt);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const auto maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const auto maxSmemPerBlock = deviceProperties->sharedMemPerBlock;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n scalarType, \"transducer_loss_cuda_forward\", ([&] {\n // resolve accumulation type\n using acc_t = at::acc_type;\n auto accType = c10::CppTypeToScalarType::value;\n auto accTensorOpt = tensorOpt.dtype(accType);\n alpha = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n beta = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n\n // decide what kernel to launch based on the problem size\n // if the required SMEM size or number threads exceeds the limit, fall back to the vanilla\n // kernel.\n const auto smemSize = 2 * maxGLen * sizeof(acc_t);\n const auto optFallBack = (maxGLen > maxThreadPerBlock or smemSize > maxSmemPerBlock) ? 0\n : (opt == -1) ? 1\n : opt;\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(2, batchSize, 1);\n\n if (optFallBack == 0)\n transducer_loss_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n else if (optFallBack == 1)\n transducer_loss_batch_load_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n }));\n C10_CUDA_CHECK(cudaGetLastError());\n\n return {alpha, beta, loss};\n}\n\ntorch::Tensor transducer_loss_cuda_backward(torch::Tensor x, torch::Tensor lossGrad, torch::Tensor alpha,\n torch::Tensor beta, torch::Tensor audLen, torch::Tensor txtLen,\n torch::Tensor label, torch::Tensor batchOffset, int maxFLen, int blankIdx,\n int opt, bool fuseSoftmaxBackward, bool packedInput) {\n auto dtype = x.scalar_type();\n torch::Tensor xGrad;\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const int maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const int warpSize = deviceProperties->warpSize;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fuseSoftmaxBackward) {\n // alloc empty tensors for performance, hence need to ensure zeros are writtern to\n // don't-care region in the kernel.\n xGrad = torch::empty_like(x);\n\n // Would like each thread to work on 4 hidden units\n const int workPerThread = 4;\n // Don't want to have more than 128 threads per thread block\n const int maxThreadPerElmt = std::min(128, maxThreadPerBlock);\n const int threads = std::min(maxThreadPerElmt, std::max(warpSize, (dictSize + workPerThread - 1) / workPerThread));\n const dim3 blocks(maxGLen, maxFLen, batchSize);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_loss_cuda_backward\", ([&] {\n using vec_t = uint64_t;\n using acc_t = at::acc_type;\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\n constexpr int vecAlignment = std::alignment_of::value;\n // if all input and output tensors meet the alignment requirement\n bool memAlign = reinterpret_cast(x.data_ptr()) % vecAlignment == 0 and\n reinterpret_cast(xGrad.data_ptr()) % vecAlignment == 0;\n\n if (vectFactor > 1 and dictSize % vectFactor == 0 and memAlign) {\n transducer_loss_fused_vec_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n } else {\n transducer_loss_fused_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }\n }));\n } else {\n // for non-fused kernel, the gradients need to be writtern are very sparse, hence initialize\n // the tensor with all zeros.\n xGrad = torch::zeros_like(x);\n // don't launch more threads than needed.\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(maxFLen, batchSize);\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, \"transducer_loss_cuda_backward\", ([&] {\n using acc_t = at::acc_type;\n transducer_loss_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(),\n audLen.data_ptr(), txtLen.data_ptr(), label.data_ptr(),\n alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }));\n }\n C10_CUDA_CHECK(cudaGetLastError());\n\n return xGrad;\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/transducer/transducer_loss_kernel.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n\n#include \n\n__global__ void index_mul_2d_float_dim64(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 res, src1, src2;\n src1 = reinterpret_cast(in1)[vec_idx1];\n src2 = reinterpret_cast(in2)[vec_idx2];\n res.x = src1.x * src2.x;\n res.y = src1.y * src2.y;\n res.z = src1.z * src2.z;\n res.w = src1.w * src2.w;\n reinterpret_cast(out)[vec_idx2] = res;\n }\n}\n\n__global__ void index_mul_2d_float(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = in1[vec_idx1 + i] * in2[vec_idx2 + i];\n }\n }\n}\n\n__global__ void index_mul_2d_half(at::Half* out, const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = at::Half(static_cast(in1[vec_idx1 + i]) * static_cast(in2[vec_idx2 + i]));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_float_dim64(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_in1, src_in2, src_grad_out, dst_grad_in2;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_out.x * src_in2.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_out.y * src_in2.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_out.z * src_in2.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_out.w * src_in2.w);\n dst_grad_in2.x = src_grad_out.x * src_in1.x;\n dst_grad_in2.y = src_grad_out.y * src_in1.y;\n dst_grad_in2.z = src_grad_out.z * src_in1.z;\n dst_grad_in2.w = src_grad_out.w * src_in1.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_float(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size,\n const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_in2[vec_idx2 + i] = src_grad_out * src_in1;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_out * src_in2);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_half(at::Half* grad_in1, at::Half* grad_in2, const at::Half* grad_out,\n const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_out * src_in1);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_out * src_in2));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float_dim64(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_grad_grad_in1, src_in1, src_grad_grad_in2, src_in2, src_grad_out;\n float4 dst_grad_grad_out, dst_grad_in2;\n src_grad_grad_in1 = reinterpret_cast(grad_grad_in1)[vec_idx1];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_grad_grad_in2 = reinterpret_cast(grad_grad_in2)[vec_idx2];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n dst_grad_grad_out.x = src_grad_grad_in1.x * src_in2.x + src_grad_grad_in2.x * src_in1.x;\n dst_grad_grad_out.y = src_grad_grad_in1.y * src_in2.y + src_grad_grad_in2.y * src_in1.y;\n dst_grad_grad_out.z = src_grad_grad_in1.z * src_in2.z + src_grad_grad_in2.z * src_in1.z;\n dst_grad_grad_out.w = src_grad_grad_in1.w * src_in2.w + src_grad_grad_in2.w * src_in1.w;\n reinterpret_cast(grad_grad_out)[vec_idx2] = dst_grad_grad_out;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_grad_in2.x * src_grad_out.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_grad_in2.y * src_grad_out.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_grad_in2.z * src_grad_out.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_grad_in2.w * src_grad_out.w);\n dst_grad_in2.x = src_grad_grad_in1.x * src_grad_out.x;\n dst_grad_in2.y = src_grad_grad_in1.y * src_grad_out.y;\n dst_grad_in2.z = src_grad_grad_in1.z * src_grad_out.z;\n dst_grad_in2.w = src_grad_grad_in1.w * src_grad_out.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = grad_grad_in1[vec_idx1 + i];\n float src_grad_grad_in2 = grad_grad_in2[vec_idx2 + i];\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_grad_out[vec_idx2 + i] = src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1;\n grad_in2[vec_idx2 + i] = src_grad_grad_in1 * src_grad_out;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_grad_in2 * src_grad_out);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_half(at::Half* grad_grad_out, at::Half* grad_in1, at::Half* grad_in2,\n const at::Half* grad_out, const at::Half* grad_grad_in1,\n const at::Half* grad_grad_in2, const at::Half* in1, const at::Half* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = static_cast(grad_grad_in1[vec_idx1 + i]);\n float src_grad_grad_in2 = static_cast(grad_grad_in2[vec_idx2 + i]);\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_grad_out[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_grad_out);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_grad_in2 * src_grad_out));\n }\n }\n}\n\nvoid index_mul_2d_float_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float_dim64<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_float_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float_dim64<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size);\n\n AT_CUDA_CHECK(cudaGetLastError());\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n}\n\nvoid index_mul_2d_float_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float_dim64<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_half_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_half<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size,\n fea_dim);\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_half_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_half<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n}\n\nvoid index_mul_2d_half_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_half<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/index_mul_2d/index_mul_2d_cuda_kernel.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n/*\n * This file implements a crude but effective mechanism for copying data between tenors owned by different ranks\n * on the same machine using cudaMemcpyAsync peer-to-peer transfers.\n */\n\nnamespace {\n\n__global__ void AddDelay_kernel(const int delay, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay);\n *counter = new_counter;\n }\n}\n\nclass NcclCommWrapper {\n private:\n ncclComm_t comm;\n int rank, world_size;\n\n ncclDataType_t get_nccl_type(at::Tensor input) {\n switch (input.scalar_type()) {\n case at::ScalarType::Half:\n return ncclFloat16;\n case at::ScalarType::Float:\n return ncclFloat32;\n case at::ScalarType::Double:\n return ncclFloat64;\n case at::ScalarType::Byte:\n return ncclUint8;\n case at::ScalarType::Char:\n return ncclInt8;\n case at::ScalarType::Int:\n return ncclInt32;\n case at::ScalarType::Long:\n return ncclInt64;\n case at::ScalarType::BFloat16:\n return ncclBfloat16;\n default:\n assert(false);\n }\n }\n\n public:\n NcclCommWrapper() {\n memset(&comm, 0, sizeof(ncclComm_t));\n rank = 0;\n world_size = 0;\n }\n NcclCommWrapper(ncclUniqueId id, int my_rank, int num_ranks) {\n ncclCommInitRank(&comm, num_ranks, id, my_rank);\n rank = my_rank;\n world_size = num_ranks;\n }\n\n ~NcclCommWrapper() {\n printf(\"ncclCommDestroy()\\n\");\n ncclCommDestroy(comm);\n }\n\n void left_right_halo_exchange_inplace(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n auto stream = at::cuda::getCurrentCUDAStream();\n ncclGroupStart();\n ncclDataType_t ncclType = get_nccl_type(left_output_halo);\n bool left_zero = (left_rank < 0);\n bool right_zero = (right_rank < 0);\n size_t left_n = torch::numel(left_output_halo);\n size_t right_n = torch::numel(right_output_halo);\n assert(left_n > 0 && left_n == right_n);\n if (left_zero) {\n left_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, left_output_halo.scalar_type(),\n \"left_halo_exch\", [&]() {\n // send left (to my_rank - 1)\n ncclSend(left_output_halo.data_ptr(), left_n, ncclType, left_rank, comm, stream);\n // receive left (from my_rank - 1)\n ncclRecv(left_input_halo.data_ptr(), right_n, ncclType, left_rank, comm, stream);\n });\n }\n if (right_zero) {\n right_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, right_output_halo.scalar_type(),\n \"right_halo_exch\", [&]() {\n // send right (to my_rank + 1 )\n ncclSend(right_output_halo.data_ptr(), right_n, ncclType, right_rank, comm, stream);\n // receive right (from my_rank + 1)\n ncclRecv(right_input_halo.data_ptr(), left_n, ncclType, right_rank, comm, stream);\n });\n }\n ncclGroupEnd();\n }\n\n std::vector left_right_halo_exchange(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n // after halo exchange:\n // left_output_halo of rank+1 ends up in right_input_halo of rank\n // right_output_halo of rank-1 ends up in left_input_halo of rank\n auto right_input_halo = torch::empty_like(left_output_halo);\n auto left_input_halo = torch::empty_like(right_output_halo);\n left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo, left_input_halo,\n right_input_halo);\n return {left_input_halo, right_input_halo};\n }\n};\n\nclass ManagedObjects {\n public:\n ManagedObjects() {}\n ~ManagedObjects() {\n for (auto it = _nccl_comms.begin(); it != _nccl_comms.end(); ++it) {\n delete *it;\n }\n }\n\n int add_comm(NcclCommWrapper* comm) {\n int handle = _nccl_comms.size();\n _nccl_comms.push_back(comm);\n return handle;\n }\n\n NcclCommWrapper& get_comm(int handle) {\n assert(handle >= 0 && handle < _nccl_comms.size());\n return *_nccl_comms[handle];\n }\n\n private:\n std::vector _nccl_comms;\n};\nclass ManagedObjects mo;\n\n} // end anonymous namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace nccl_p2p {\n\nat::Tensor get_unique_nccl_id(int n) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n auto id_tensor = torch::empty({n, (int)sizeof(ncclUniqueId)},\n torch::dtype(torch::kUInt8).device(torch::kCPU).requires_grad(false));\n auto id_ptr = id_tensor.data_ptr();\n size_t offset = 0;\n for (int i = 0; i < n; ++i) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n memcpy(id_ptr + offset, &id, sizeof(ncclUniqueId));\n offset += sizeof(ncclUniqueId);\n }\n return id_tensor;\n}\n\nint init_nccl_comm(at::Tensor unique_nccl_id, int my_rank, int num_ranks) {\n ncclUniqueId id;\n auto unique_nccl_id_ptr = unique_nccl_id.data_ptr();\n memcpy(&id, unique_nccl_id_ptr, sizeof(ncclUniqueId));\n NcclCommWrapper* comm = new NcclCommWrapper(id, my_rank, num_ranks);\n int handle = mo.add_comm(comm);\n comm = 0L;\n return handle;\n}\n\nvoid left_right_halo_exchange_inplace(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo,\n left_input_halo, right_input_halo);\n}\n\nstd::vector left_right_halo_exchange(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange(left_rank, right_rank, left_output_halo, right_output_halo);\n}\n\nvoid add_delay(int delay) {\n auto stream = at::cuda::getCurrentCUDAStream();\n auto t = torch::empty({1}, torch::dtype(torch::kInt32).device(torch::kCUDA));\n AddDelay_kernel<<<1, 1, 0, stream>>>(delay, t.data_ptr());\n}\n\n} // namespace nccl_p2p\n} // namespace contrib\n} // namespace apex\n\n", "output": "To be filled or generated by LLM teacher model", "source": "/kaggle/working/apex/apex/contrib/csrc/nccl_p2p/nccl_p2p_cuda.cu"} {"instruction": "Analyze CUDA kernel architecture", "input": "\nExplain this CUDA kernel at architectural level:\n- memory model (global/shared/register)\n- parallel pattern\n- optimization strategy\n- bottlenecks\n\nKernel:\n/**\n * From PyTorch:\n *\n * Copyright (c) 2016- Facebook, Inc (Adam Paszke)\n * Copyright (c) 2014- Facebook, Inc (Soumith Chintala)\n * Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n * Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu)\n * Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n * Copyright (c) 2011-2013 NYU (Clement Farabet)\n * Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n * Copyright (c) 2006 Idiap Research Institute (Samy Bengio)\n * Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n *\n * From Caffe2:\n *\n * Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n *\n * All contributions by Facebook:\n * Copyright (c) 2016 Facebook Inc.\n *\n * All contributions by Google:\n * Copyright (c) 2015 Google Inc.\n * All rights reserved.\n *\n * All contributions by Yangqing Jia:\n * Copyright (c) 2015 Yangqing Jia\n * All rights reserved.\n *\n * All contributions from Caffe:\n * Copyright(c) 2013, 2014, 2015, the respective contributors\n * All rights reserved.\n *\n * All other contributions:\n * Copyright(c) 2015, 2016 the respective contributors\n * All rights reserved.\n *\n * Caffe2 uses a copyright model similar to Caffe: each contributor holds\n * copyright over their contributions to Caffe2. The project versioning records\n * all such contribution and copyright details. If a contributor wants to further\n * mark their specific copyright on a particular contribution, they should\n * indicate their copyright solely in the commit message of the change when it is\n * committed.\n *\n * All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n *\n * 1. Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n *\n * 2. Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n *\n * 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n * and IDIAP Research Institute nor the names of its contributors may be\n * used to endorse or promote products derived from this software without\n * specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n * POSSIBILITY OF SUCH DAMAGE.\n */\n#include \n#include \n#include \n\n#include \n\n#include \"type_shim.h\"\n\n#define ALIGN_BYTES 16\n\nusing Tensor = at::Tensor;\nusing TensorList = at::TensorList;\nusing ScalarType = at::ScalarType;\nusing at::acc_type;\n\ntemplate \nstruct LogSoftMaxForwardEpilogue {\n __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_input, AccumT sum)\n : logsum(max_input + std::log(sum)) {}\n\n __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_log_sum_exp) : logsum(max_log_sum_exp) {}\n\n __device__ __forceinline__ OutT operator()(T input) const { return static_cast(input - logsum); }\n\n const AccumT logsum;\n};\n\ntemplate \nstruct LogSoftMaxBackwardEpilogue {\n __device__ __forceinline__ LogSoftMaxBackwardEpilogue(AccumT sum) : sum(sum) {}\n\n __device__ __forceinline__ T operator()(OutT gradOutput, OutT output) const {\n return static_cast(gradOutput - std::exp(static_cast(output)) * sum);\n }\n\n const AccumT sum;\n};\n\nconst int max_threads = 1024;\n\ninline dim3 SoftMax_getBlockSize(int ILP, uint64_t dim_size) {\n uint64_t block_size = 1;\n uint64_t max_block_size = std::min(dim_size / ILP, static_cast(max_threads));\n while (block_size < (max_block_size / 2)) block_size *= 2;\n // Launch at least a single warp - the kernel assumes that.\n block_size = std::max(block_size, static_cast(32));\n return dim3(block_size);\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Regular kernel (fast when dim_size is large; requires inner_size == 1)\n////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct MaxFloat {\n __device__ __forceinline__ AccumT operator()(AccumT max, T v) const { return ::max(max, (AccumT)v); }\n};\n\ntemplate \nstruct AddFloat {\n __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const { return sum + v; }\n};\n\ntemplate \nstruct SumExpFloat {\n __device__ __forceinline__ SumExpFloat(AccumT v) : max_k(v) {}\n\n __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const { return sum + std::exp(v - max_k); }\n\n const AccumT max_k;\n};\n\ntemplate