34 lines
1.0 KiB
Rust
34 lines
1.0 KiB
Rust
use crate::ir::{DataType, IRNode, IRNodeType, NodeId, Shape, SparseFormat};
|
|
|
|
/// Sparse tensor operations for the Unified IR
|
|
pub struct SparseOps;
|
|
|
|
impl SparseOps {
|
|
pub fn create_sparse_matmul(
|
|
id: NodeId,
|
|
inputs: Vec<NodeId>,
|
|
format: SparseFormat,
|
|
output_shape: Shape,
|
|
) -> IRNode {
|
|
IRNode::new(
|
|
id,
|
|
IRNodeType::SparseMatMul { format },
|
|
inputs,
|
|
vec![DataType::F32],
|
|
vec![output_shape],
|
|
)
|
|
}
|
|
|
|
pub fn create_csr_matmul(id: NodeId, inputs: Vec<NodeId>, output_shape: Shape) -> IRNode {
|
|
Self::create_sparse_matmul(id, inputs, SparseFormat::CSR, output_shape)
|
|
}
|
|
|
|
pub fn create_coo_matmul(id: NodeId, inputs: Vec<NodeId>, output_shape: Shape) -> IRNode {
|
|
Self::create_sparse_matmul(id, inputs, SparseFormat::COO, output_shape)
|
|
}
|
|
|
|
pub fn create_csc_matmul(id: NodeId, inputs: Vec<NodeId>, output_shape: Shape) -> IRNode {
|
|
Self::create_sparse_matmul(id, inputs, SparseFormat::CSC, output_shape)
|
|
}
|
|
}
|