Function af_matmul

src/api/c/blas.cpp:235–315

}

af_err af_matmul(af_array *out, const af_array lhs, const af_array rhs,
                 const af_mat_prop optLhs, const af_mat_prop optRhs) {
    try {
        const ArrayInfo &lhsInfo = getInfo(lhs, false);
        const ArrayInfo &rhsInfo = getInfo(rhs, true);

        if (lhsInfo.isSparse()) {
            return af_sparse_matmul(out, lhs, rhs, optLhs, optRhs);
        }

        const int aRowDim = (optLhs == AF_MAT_NONE) ? 0 : 1;
        const int bColDim = (optRhs == AF_MAT_NONE) ? 1 : 0;

        const af::dim4 &lDims = lhsInfo.dims();
        const af::dim4 &rDims = rhsInfo.dims();
        const int M           = lDims[aRowDim];
        const int N           = rDims[bColDim];

        const dim_t d2       = std::max(lDims[2], rDims[2]);
        const dim_t d3       = std::max(lDims[3], rDims[3]);
        const af::dim4 oDims = af::dim4(M, N, d2, d3);

        af_dtype lhs_type = lhsInfo.getType();

        af_array gemm_out      = 0;
        af_dtype gemm_out_type = (lhs_type != s8) ? lhs_type : f32;
        AF_CHECK(af_create_handle(&gemm_out, oDims.ndims(), oDims.get(),
                                  gemm_out_type));

        switch (lhs_type) {
            case f16: {
                static const half alpha(1.0f);
                static const half beta(0.0f);
                AF_CHECK(af_gemm(&gemm_out, optLhs, optRhs, &alpha, lhs, rhs,
                                 &beta));
                break;
            }
            case f32: {
                float alpha = 1.f;
                float beta  = 0.f;
                AF_CHECK(af_gemm(&gemm_out, optLhs, optRhs, &alpha, lhs, rhs,
                                 &beta));
                break;
            }
            case c32: {
                cfloat alpha{1.f, 0.f};
                cfloat beta{0.f, 0.f};

                AF_CHECK(af_gemm(&gemm_out, optLhs, optRhs, &alpha, lhs, rhs,
                                 &beta));
                break;
            }
            case f64: {
                double alpha = 1.0;
                double beta  = 0.0;
                AF_CHECK(af_gemm(&gemm_out, optLhs, optRhs, &alpha, lhs, rhs,
                                 &beta));
                break;