| 132 | } |
| 133 | |
| 134 | af_err af_gemm(af_array *out, const af_mat_prop optLhs, |
| 135 | const af_mat_prop optRhs, const void *alpha, const af_array lhs, |
| 136 | const af_array rhs, const void *beta) { |
| 137 | try { |
| 138 | const ArrayInfo &lhsInfo = getInfo(lhs, false); |
| 139 | const ArrayInfo &rhsInfo = getInfo(rhs, true); |
| 140 | |
| 141 | af_dtype lhs_type = lhsInfo.getType(); |
| 142 | af_dtype rhs_type = rhsInfo.getType(); |
| 143 | |
| 144 | if (!(optLhs == AF_MAT_NONE || optLhs == AF_MAT_TRANS || |
| 145 | optLhs == AF_MAT_CTRANS)) { |
| 146 | AF_ERROR("Using this property is not yet supported in matmul", |
| 147 | AF_ERR_NOT_SUPPORTED); |
| 148 | } |
| 149 | |
| 150 | if (!(optRhs == AF_MAT_NONE || optRhs == AF_MAT_TRANS || |
| 151 | optRhs == AF_MAT_CTRANS)) { |
| 152 | AF_ERROR("Using this property is not yet supported in matmul", |
| 153 | AF_ERR_NOT_SUPPORTED); |
| 154 | } |
| 155 | |
| 156 | af::dim4 lDims = lhsInfo.dims(); |
| 157 | af::dim4 rDims = rhsInfo.dims(); |
| 158 | |
| 159 | if (lDims.ndims() > 2 && rDims.ndims() > 2) { |
| 160 | DIM_ASSERT(3, lDims.ndims() == rDims.ndims()); |
| 161 | if (lDims[2] != rDims[2] && lDims[2] != 1 && rDims[2] != 1) { |
| 162 | AF_ERROR("Batch size mismatch along dimension 2", AF_ERR_BATCH); |
| 163 | } |
| 164 | if (lDims[3] != rDims[3] && lDims[3] != 1 && rDims[3] != 1) { |
| 165 | AF_ERROR("Batch size mismatch along dimension 3", AF_ERR_BATCH); |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | TYPE_ASSERT(lhs_type == rhs_type); |
| 170 | |
| 171 | int aColDim = (optLhs == AF_MAT_NONE) ? 1 : 0; |
| 172 | int bRowDim = (optRhs == AF_MAT_NONE) ? 0 : 1; |
| 173 | |
| 174 | DIM_ASSERT(1, lhsInfo.dims()[aColDim] == rhsInfo.dims()[bRowDim]); |
| 175 | |
| 176 | // Assume that *out is either initialized to null or an actual af_array |
| 177 | // Otherwise, this function has undefined behavior |
| 178 | af_array output = 0; |
| 179 | if (*out) { |
| 180 | output = *out; |
| 181 | } else { |
| 182 | af_dtype out_type = (lhs_type != s8) ? lhs_type : f32; |
| 183 | |
| 184 | const int aRowDim = (optLhs == AF_MAT_NONE) ? 0 : 1; |
| 185 | const int bColDim = (optRhs == AF_MAT_NONE) ? 1 : 0; |
| 186 | const int M = lDims[aRowDim]; |
| 187 | const int N = rDims[bColDim]; |
| 188 | const dim_t d2 = std::max(lDims[2], rDims[2]); |
| 189 | const dim_t d3 = std::max(lDims[3], rDims[3]); |
| 190 | const af::dim4 oDims = af::dim4(M, N, d2, d3); |
| 191 | AF_CHECK(af_create_handle(&output, lhsInfo.ndims(), oDims.get(), |