Create fused_ops.cu

core/kernels/fused_ops.cu

@@ -0,0 +1,137 @@
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+#include <mma.h>  // Tensor Core WMMA API
+
+using namespace nvcuda;
+
+// Define FP16 type for Tensor Cores
+using half_t = __half;
+
+// Thread block and warp sizes
+#define BLOCK_SIZE 32
+#define WARP_SIZE 32
+#define WMMA_M 16
+#define WMMA_N 16
+#define WMMA_K 16
+
+// Fused sparse GEMM + ReLU kernel
+__global__ void fused_sparse_gemm_relu_kernel(
+    const half_t* __restrict__ input,      // Input tensor [batch_size, in_features]
+    const half_t* __restrict__ weight,     // Weight tensor [out_features, in_features]
+    const half_t* __restrict__ mask,       // Sparsity mask [out_features, in_features]
+    half_t* __restrict__ output,           // Output tensor [batch_size, out_features]
+    const half_t* __restrict__ bias,       // Bias tensor [out_features]
+    int batch_size, int in_features, int out_features)
+{
+    // Shared memory for WMMA fragments
+    __shared__ half_t shmem_input[BLOCK_SIZE * WMMA_K];
+    __shared__ half_t shmem_weight[WMMA_M * WMMA_K];
+
+    // WMMA fragments
+    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half_t, wmma::row_major> a_frag;
+    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half_t, wmma::col_major> b_frag;
+    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, half_t> c_frag;
+
+    // Thread indices
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+    int tx = threadIdx.x;
+    int ty = threadIdx.y;
+
+    // Global indices
+    int row = by * WMMA_M + ty;  // Output row
+    int col = bx * WMMA_N + tx;  // Output col
+
+    // Compute tile offsets
+    int batch_offset = blockIdx.z * in_features;  // Batch dimension
+
+    // Initialize accumulator
+    wmma::fill_fragment(c_frag, __float2half(0.0f));
+
+    // Loop over K dimension (in_features) in WMMA tiles
+    for (int k = 0; k < in_features; k += WMMA_K) {
+        // Load input into shared memory
+        if (ty < WMMA_K && row < batch_size) {
+            shmem_input[ty * BLOCK_SIZE + tx] = input[batch_offset + row * in_features + k + tx];
+        }
+
+        // Load sparse weight into shared memory (apply mask)
+        if (ty < WMMA_M && k + tx < in_features && row < out_features) {
+            half_t w = weight[row * in_features + k + tx];
+            half_t m = mask[row * in_features + k + tx];
+            shmem_weight[ty * WMMA_K + tx] = __hmul(w, m);  // Apply sparsity mask
+        }
+
+        __syncthreads();
+
+        // Load WMMA fragments from shared memory
+        wmma::load_matrix_sync(a_frag, shmem_input, BLOCK_SIZE);
+        wmma::load_matrix_sync(b_frag, shmem_weight, WMMA_K);
+
+        // Perform Tensor Core matrix multiply-accumulate
+        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
+
+        __syncthreads();
+    }
+
+    // Store result with ReLU and bias
+    if (row < batch_size && col < out_features) {
+        half_t result = c_frag.x[ty * WMMA_N + tx];
+        result = __hadd(result, bias[col]);  // Add bias
+        output[row * out_features + col] = __hgt(result, __float2half(0.0f)) ? result : __float2half(0.0f);  // ReLU
+    }
+}
+
+// PyTorch binding
+torch::Tensor fused_sparse_gemm_relu(
+    torch::Tensor input,    // [batch_size, in_features]
+    torch::Tensor weight,   // [out_features, in_features]
+    torch::Tensor mask,     // [out_features, in_features]
+    torch::Tensor bias)     // [out_features]
+{
+    // Ensure inputs are FP16 and on CUDA
+    TORCH_CHECK(input.dtype() == torch::kFloat16, "Input must be FP16");
+    TORCH_CHECK(weight.dtype() == torch::kFloat16, "Weight must be FP16");
+    TORCH_CHECK(mask.dtype() == torch::kFloat16, "Mask must be FP16");
+    TORCH_CHECK(bias.dtype() == torch::kFloat16, "Bias must be FP16");
+    TORCH_CHECK(input.is_cuda(), "Input must be on CUDA");
+    TORCH_CHECK(weight.is_cuda(), "Weight must be on CUDA");
+    TORCH_CHECK(mask.is_cuda(), "Mask must be on CUDA");
+    TORCH_CHECK(bias.is_cuda(), "Bias must be on CUDA");
+
+    // Dimensions
+    int batch_size = input.size(0);
+    int in_features = input.size(1);
+    int out_features = weight.size(0);
+
+    // Output tensor
+    auto output = torch::empty({batch_size, out_features}, 
+                              torch::TensorOptions().dtype(torch::kFloat16).device(input.device()));
+
+    // Grid and block dimensions
+    dim3 block(BLOCK_SIZE, WMMA_M / WARP_SIZE);
+    dim3 grid((out_features + WMMA_N - 1) / WMMA_N, 
+              (batch_size + WMMA_M - 1) / WMMA_M, 
+              batch_size);
+
+    // Launch kernel
+    fused_sparse_gemm_relu_kernel<<<grid, block>>>(
+        (half_t*)input.data_ptr(),
+        (half_t*)weight.data_ptr(),
+        (half_t*)mask.data_ptr(),
+        (half_t*)output.data_ptr(),
+        (half_t*)bias.data_ptr(),
+        batch_size, in_features, out_features);
+
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        TORCH_CHECK(false, "CUDA error: ", cudaGetErrorString(err));
+    }
+
+    return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fused_sparse_gemm_relu", &fused_sparse_gemm_relu, "Fused sparse GEMM + ReLU with Tensor Cores");
+}