Function solve

src/backend/cpu/solve.cpp:262–323

template<typename T>
Array<T> solve(const Array<T> &a, const Array<T> &b,
               const af_mat_prop options) {
    if (options & AF_MAT_UPPER || options & AF_MAT_LOWER) {
        return triangleSolve<T>(a, b, options);
    }

#ifdef AF_USE_MKL_BATCH
    if (a.dims()[2] > 1 || a.dims()[3] > 1) {
        return generalSolveBatched(a, b, options);
    }
#endif

    const dim4 NullShape(0, 0, 0, 0);

    dim4 aDims = a.dims();
    int batchz = aDims[2];
    int batchw = aDims[3];

    int M = aDims[0];
    int N = aDims[1];
    int K = b.dims()[1];

    Array<T> A = copyArray<T>(a);

    dim4 endPadding(max(M, N) - b.dims()[0], K - b.dims()[1], 0, 0);
    Array<T> B = (endPadding == NullShape
                      ? copyArray(b)
                      : padArrayBorders(b, NullShape, endPadding, AF_PAD_ZERO));

    for (int i = 0; i < batchw; i++) {
        for (int j = 0; j < batchz; j++) {
            Param<T> pA(A.get() + A.strides()[2] * j + A.strides()[3] * i,
                        A.dims(), A.strides());
            Param<T> pB(B.get() + B.strides()[2] * j + B.strides()[3] * i,
                        B.dims(), B.strides());
            if (M == N) {
                Array<int> pivot = createEmptyArray<int>(dim4(N, 1, 1));

                auto func = [](Param<T> A, Param<T> B, Param<int> pivot, int N,
                               int K) {
                    gesv_func<T>()(AF_LAPACK_COL_MAJOR, N, K, A.get(),
                                   A.strides(1), pivot.get(), B.get(),
                                   B.strides(1));
                };
                getQueue().enqueue(func, pA, pB, pivot, N, K);
            } else {
                auto func = [=](Param<T> A, Param<T> B, int M, int N, int K) {
                    int sM = A.dims(0);
                    int sN = A.dims(1);

                    gels_func<T>()(AF_LAPACK_COL_MAJOR, 'N', M, N, K, A.get(),
                                   A.strides(1), B.get(), max(sM, sN));
                };
                getQueue().enqueue(func, pA, pB, M, N, K);
            }
        }
    }