BlockSmith: Ultra-dense Block Transformation for Efficient Sparse Matrix Multiplication Using Tensor Cores

This repository contains the official implementation of BlockSmith, including GPU preprocessing and Tensor Core SpMM evaluation.

Installation

1. Create Conda Environment

conda create -n blocksmith python=3.9
conda activate blocksmith

2. Install Dependencies

This project requires CUDA, a CUDA-capable NVIDIA GPU with Tensor Core support, and a host compiler supported by the installed CUDA toolkit.

conda install -c conda-forge "gcc<12" "gxx<12"

3. Compile Project Components

Set the CUDA architecture for your GPU before building.

SM=89 DEV_ID=0 scripts/build.sh

The build script compiles:

preprocess/programs/preprocess_main_cuda
preprocess/programs/reorder_main_cuda
build/spmm/tcspmm_fp16
build/spmm/tcspmm_tf32_n{32,64,128,256}

Dataset Preparation

The benchmark matrix list is provided in:

data/matrices_manifest.csv

Download and extract the SuiteSparse matrices:

scripts/download_datasets.py

Matrices are stored in:

datasets/<matrix>/<matrix>.mtx

BlockSmith Workflow

The standard workflow is:

1. Download the benchmark matrices.
2. Run GPU preprocessing to generate reordered matrices and RC-BCSR files.
3. Run SpMM kernel evaluation.
4. Summarize results.

Step 1: GPU Preprocessing

Run the GPU preprocessing pipeline:

scripts/run_preprocess.py

Outputs are written to:

outputs/preprocess/preprocess_results.csv
outputs/preprocess/reordered_mtx/<matrix>/<matrix>.mtx
outputs/preprocess/rcbcsr/<matrix>/<matrix>.txt
outputs/preprocess/reordered_blocks_log.csv

The preprocessing report includes the GPU reorder timing used for overhead evaluation:

preprocess_ms

Generate the preprocessing-overhead summary:

scripts/summarize_preprocess.py

The summary files are:

outputs/preprocess/blocksmith_reorder_overhead_n256.csv
outputs/preprocess/blocksmith_reorder_overhead_n256_sorted_by_overhead.csv

The reference CSV files in results_reference/ were produced on an NVIDIA RTX 4090 GPU.

Step 2: SpMM Kernel Evaluation

Run the FP16 SpMM evaluation:

scripts/run_spmm_fp16.py

The output CSV is:

outputs/spmm/blocksmith_fp16_n32.csv

Run TF32 SpMM evaluation for multiple dense feature dimensions:

scripts/run_spmm_tf32.py --n 32 --n 64 --n 128 --n 256

The output CSVs are:

outputs/spmm/blocksmith_tf32_n32.csv
outputs/spmm/blocksmith_tf32_n64.csv
outputs/spmm/blocksmith_tf32_n128.csv
outputs/spmm/blocksmith_tf32_n256.csv

Each SpMM CSV reports the kernel timings and the best observed BlockSmith kernel for each matrix.