| 69 | } // namespace |
| 70 | |
| 71 | af_err af_sparse_matmul(af_array *out, const af_array lhs, const af_array rhs, |
| 72 | const af_mat_prop optLhs, const af_mat_prop optRhs) { |
| 73 | try { |
| 74 | const SparseArrayBase lhsBase = getSparseArrayBase(lhs); |
| 75 | const ArrayInfo &rhsInfo = getInfo(rhs); |
| 76 | |
| 77 | ARG_ASSERT(2, |
| 78 | lhsBase.isSparse() == true && rhsInfo.isSparse() == false); |
| 79 | |
| 80 | af_dtype lhs_type = lhsBase.getType(); |
| 81 | af_dtype rhs_type = rhsInfo.getType(); |
| 82 | |
| 83 | ARG_ASSERT(1, lhsBase.getStorage() == AF_STORAGE_CSR); |
| 84 | |
| 85 | if (!(optLhs == AF_MAT_NONE || optLhs == AF_MAT_TRANS || |
| 86 | optLhs == AF_MAT_CTRANS)) { // Note the ! operator. |
| 87 | AF_ERROR( |
| 88 | "Using this property is not yet supported in sparse matmul", |
| 89 | AF_ERR_NOT_SUPPORTED); |
| 90 | } |
| 91 | |
| 92 | // No transpose options for RHS |
| 93 | if (optRhs != AF_MAT_NONE) { |
| 94 | AF_ERROR("Using this property is not yet supported in matmul", |
| 95 | AF_ERR_NOT_SUPPORTED); |
| 96 | } |
| 97 | |
| 98 | if (rhsInfo.ndims() > 2) { |
| 99 | AF_ERROR("Sparse matmul can not be used in batch mode", |
| 100 | AF_ERR_BATCH); |
| 101 | } |
| 102 | |
| 103 | TYPE_ASSERT(lhs_type == rhs_type); |
| 104 | |
| 105 | af::dim4 ldims = lhsBase.dims(); |
| 106 | int lColDim = (optLhs == AF_MAT_NONE) ? 1 : 0; |
| 107 | int rRowDim = (optRhs == AF_MAT_NONE) ? 0 : 1; |
| 108 | |
| 109 | DIM_ASSERT(1, ldims[lColDim] == rhsInfo.dims()[rRowDim]); |
| 110 | |
| 111 | af_array output = 0; |
| 112 | switch (lhs_type) { |
| 113 | case f32: |
| 114 | output = sparseMatmul<float>(lhs, rhs, optLhs, optRhs); |
| 115 | break; |
| 116 | case c32: |
| 117 | output = sparseMatmul<cfloat>(lhs, rhs, optLhs, optRhs); |
| 118 | break; |
| 119 | case f64: |
| 120 | output = sparseMatmul<double>(lhs, rhs, optLhs, optRhs); |
| 121 | break; |
| 122 | case c64: |
| 123 | output = sparseMatmul<cdouble>(lhs, rhs, optLhs, optRhs); |
| 124 | break; |
| 125 | default: TYPE_ERROR(1, lhs_type); |
| 126 | } |
| 127 | std::swap(*out, output); |
| 128 | } |