CSR和COO实现spgemm

1) COO实现spgemm 

#include <iostream>
#include <vector>
#include <unordered_map>

// 定义稀疏矩阵的元素
struct Element {
    int row, col;
    double value;
};

// 稀疏矩阵乘法函数
std::vector<Element> spgemm(const std::vector<Element>& A, const std::vector<Element>& B, int rowsA, int colsB) {
    //创建可自由缩放的二维矩阵数组
    std::unordered_map<int, std::unordered_map<int, double>> result;

    std::unordered_map<int, std::vector<Element>> B_col;
    for (const auto& elem : B) {
        B_col[elem.col].push_back(elem);// 将B矩阵转换为列主序, 为下面的判断做准备
    }

    // 进行乘法运算
    for (const auto& a : A) {
        //判断矩阵A的列在矩阵B中是否有对应的行相匹配
        if (B_col.find(a.col) != B_col.end()) {
            //a是向量A中的一个数据类型为element的元素, a.col代表元素element.col
            //B_col[a.col]: 确保矩阵B的行与矩阵A的列相匹配
            for (const auto& b : for (const auto& b : B_col[a.col]) {
                result[a.row][b.col] += a.value * b.value;
            }) {
                result[a.row][b.col] += a.value * b.value;
            }
        }
    }

    // 将结果转换为COO格式
    std::vector<Element> C;
    for (const auto& row : result) {
        for (const auto& col : row.second) {
            C.push_back({row.first, col.first, col.second});
        }
    }
    return C;
}

int main() {
    // 示例稀疏矩阵A和B
    std::vector<Element> A = {{0, 0, 1.0}, {0, 1, 2.0}, {1, 0, 3.0}};
    std::vector<Element> B = {{0, 0, 4.0}, {1, 0, 5.0}, {1, 1, 6.0}};

    // 计算A * B
    std::vector<Element> C = spgemm(A, B, 2, 2);

    // 输出结果
    for (const auto& elem : C) {
        std::cout << "C(" << elem.row << ", " << elem.col << ") = " << elem.value << std::endl;
    }

    return 0;
}

2) CSR实现spgemm

#include <iostream>
#include <vector>
#include <unordered_map>
#include <stdexcept>
using namespace std;

// CSR格式的稀疏矩阵
struct CSRMatrix {
	std::vector<int> row_ptr;  // 每行的起始索引
	std::vector<int> col_idx;  // 非零元素的列索引
	std::vector<double> values;  // 非零元素的值
	int rows, cols;  // 矩阵的行数和列数
};

// SpGEMM算法实现
CSRMatrix spgemm(const CSRMatrix& A, const CSRMatrix& B) {
	// 检查矩阵整体维度是否匹配
	if (A.cols != B.rows) {
		throw std::invalid_argument("Matrix dimensions do not match for multiplication.");
	}

	CSRMatrix C;
	//初始化C矩阵的rows和cols
	C.rows = A.rows;
	C.cols = B.cols;
	//初始化c.row_ptr, 并将元素值都设为0
	C.row_ptr.resize(C.rows + 1, 0);

	// //定义元素个数为C.rows, 元素类型为字典的向量, 临时存储每行的非零元素
	// 实现了可以自动扩展大小的二级数组的效果
	std::vector<std::unordered_map<int, double>> temp(C.rows);

	// 遍历矩阵A的每一行
	for (int i = 0; i < A.rows; ++i) {
		// 遍历A的当前行的每一个非零元素
		//由于CSR存储格式的特征, row_ptr[i+1]-row_ptr[i]代表第i行元素的个数
		for (int j = A.row_ptr[i]; j < A.row_ptr[i + 1]; ++j) {
			int a_col = A.col_idx[j];  // A的列索引
			double a_val = A.values[j];  // A的值
			//B.row_ptr[a_col]: 找到与A的列对应的矩阵B的行
			for (int k = B.row_ptr[a_col]; k < B.row_ptr[a_col + 1]; ++k) {
				int b_col = B.col_idx[k];  // B的列索引, 与矩阵A中某一行相乘的矩阵B中的某一个列号
				double b_val = B.values[k];  // B的值
				temp[i][b_col] += a_val * b_val;  // 累加结果
			}
		}
	}

	// 将临时存储的结果转换为CSR格式
	for (int i = 0; i < C.rows; ++i) {
		for (const auto& pair : temp[i]) {
			C.col_idx.push_back(pair.first);  // 列索引
			C.values.push_back(pair.second);  // 值
		}
		C.row_ptr[i + 1] = C.col_idx.size();  // 更新行指针
	}

	return C;
}

int main() {
	
	// 示例矩阵A的初始化
	CSRMatrix A = {
		{0, 2, 4},  // 行指针
		{0, 1, 0, 2},  // 列索引
		{1.0, 2.0, 3.0, 4.0},  // 值
		2, 3  // 行数和列数
	};

	// 示例矩阵B的初始化
	CSRMatrix B = {
		{0, 1, 3, 4},  // 行指针
		{0, 1, 2, 2},  // 列索引
		{5.0, 6.0, 7.0, 8.0},  // 值
		3, 3  // 行数和列数
	};

	// 计算矩阵C = A * B
	CSRMatrix C = spgemm(A, B);

	// 输出结果矩阵C
	std::cout << "C.row_ptr: ";
	for (int val : C.row_ptr) std::cout << val << " ";
	std::cout << "\nC.col_idx: ";
	for (int val : C.col_idx) std::cout << val << " ";
	std::cout << "\nC.values: ";
	for (double val : C.values) std::cout << val << " ";
	std::cout << std::endl;
	return 0;
}