| 99 | |
| 100 | template<typename Ti, typename To> |
| 101 | void gemm(Array<To> &out, af_mat_prop optLhs, af_mat_prop optRhs, |
| 102 | const To *alpha, const Array<Ti> &lhs, const Array<Ti> &rhs, |
| 103 | const To *beta) { |
| 104 | const auto lOpts = toBlasTranspose(optLhs); |
| 105 | const auto rOpts = toBlasTranspose(optRhs); |
| 106 | |
| 107 | const auto aRowDim = (optLhs == AF_MAT_NONE) ? 0 : 1; |
| 108 | const auto aColDim = (optLhs == AF_MAT_NONE) ? 1 : 0; |
| 109 | const auto bColDim = (optRhs == AF_MAT_NONE) ? 1 : 0; |
| 110 | |
| 111 | const dim4 &lDims = lhs.dims(); |
| 112 | const dim4 &rDims = rhs.dims(); |
| 113 | const int M = lDims[aRowDim]; |
| 114 | const int N = rDims[bColDim]; |
| 115 | const int K = lDims[aColDim]; |
| 116 | const dim4 oDims = out.dims(); |
| 117 | |
| 118 | const dim4 &lStrides = lhs.strides(); |
| 119 | const dim4 &rStrides = rhs.strides(); |
| 120 | const dim4 oStrides = out.strides(); |
| 121 | |
| 122 | if (oDims.ndims() <= 2) { // if non-batched |
| 123 | if (rhs.dims()[bColDim] == 1) { |
| 124 | if constexpr (std::is_same_v<Ti, arrayfire::common::half>) { |
| 125 | // currently no half support for gemv, use gemm instead |
| 126 | gemmDispatch<Ti>(getQueue(), lOpts, rOpts, M, N, K, alpha, lhs, |
| 127 | lStrides[1], rhs, rStrides[1], beta, out, |
| 128 | oStrides[1]); |
| 129 | } else { |
| 130 | dim_t incr = |
| 131 | (optRhs == AF_MAT_NONE) ? rStrides[0] : rStrides[1]; |
| 132 | gemvDispatch<Ti>(getQueue(), lOpts, rOpts, lDims[0], lDims[1], |
| 133 | alpha, lhs, lStrides[1], rhs, incr, beta, out, |
| 134 | oStrides[0]); |
| 135 | } |
| 136 | } else { |
| 137 | gemmDispatch<Ti>(getQueue(), lOpts, rOpts, M, N, K, alpha, lhs, |
| 138 | lStrides[1], rhs, rStrides[1], beta, out, |
| 139 | oStrides[1]); |
| 140 | } |
| 141 | } else { // if batched |
| 142 | using Dt = arrayfire::oneapi::data_t<Ti>; |
| 143 | |
| 144 | int64_t batchSize = static_cast<int64_t>(oDims[2] * oDims[3]); |
| 145 | |
| 146 | bool is_l_d2_batched = (oDims[2] == lDims[2]) && lDims[2] != 1; |
| 147 | bool is_l_d3_batched = (oDims[3] == lDims[3]) && lDims[3] != 1; |
| 148 | bool is_r_d2_batched = (oDims[2] == rDims[2]) && rDims[2] != 1; |
| 149 | bool is_r_d3_batched = (oDims[3] == rDims[3]) && rDims[3] != 1; |
| 150 | |
| 151 | // MKL requires stridec >= ldc * n, which may not be true with reordered |
| 152 | // outputs if the stride is monotonic, then MKL requirements for |
| 153 | // batching can be met |
| 154 | bool canBatchMKL = isStrideMonotonic(oStrides); |
| 155 | if (canBatchMKL) { |
| 156 | sycl::buffer<Dt, 1> lhsBuf = lhs.template getBufferWithOffset<Dt>(); |
| 157 | sycl::buffer<Dt, 1> rhsBuf = rhs.template getBufferWithOffset<Dt>(); |
| 158 | sycl::buffer<Dt, 1> outBuf = out.template getBufferWithOffset<Dt>(); |
no test coverage detected