C++ ArrayInfo::dims方法代码示例

static af_err af_bitwise(af_array *out, const af_array lhs, const af_array rhs, const bool batchMode)
{
    try {
        const af_dtype type = implicit(lhs, rhs);

        ArrayInfo linfo = getInfo(lhs);
        ArrayInfo rinfo = getInfo(rhs);

        dim4 odims = getOutDims(linfo.dims(), rinfo.dims(), batchMode);

        af_array res;
        switch (type) {
        case s32: res = bitOp<int    , op>(lhs, rhs, odims); break;
        case u32: res = bitOp<uint   , op>(lhs, rhs, odims); break;
        case u8 : res = bitOp<uchar  , op>(lhs, rhs, odims); break;
        case b8 : res = bitOp<char   , op>(lhs, rhs, odims); break;
        case s64: res = bitOp<intl   , op>(lhs, rhs, odims); break;
        case u64: res = bitOp<uintl  , op>(lhs, rhs, odims); break;
        case s16: res = bitOp<short  , op>(lhs, rhs, odims); break;
        case u16: res = bitOp<ushort , op>(lhs, rhs, odims); break;
        default: TYPE_ERROR(0, type);
        }

        std::swap(*out, res);
    }
    CATCHALL;
    return AF_SUCCESS;
}

af_err af_hypot(af_array *out, const af_array lhs, const af_array rhs, const bool batchMode)
{
    try {

        const af_dtype type = implicit(lhs, rhs);

        if (type != f32 && type != f64) {
            AF_ERROR("Only floating point arrays are supported for hypot ",
                     AF_ERR_NOT_SUPPORTED);
        }

        ArrayInfo linfo = getInfo(lhs);
        ArrayInfo rinfo = getInfo(rhs);

        dim4 odims = getOutDims(linfo.dims(), rinfo.dims(), batchMode);

        af_array res;
        switch (type) {
        case f32: res = arithOp<float , af_hypot_t>(lhs, rhs, odims); break;
        case f64: res = arithOp<double, af_hypot_t>(lhs, rhs, odims); break;
        default: TYPE_ERROR(0, type);
        }

        std::swap(*out, res);
    }
    CATCHALL;
    return AF_SUCCESS;
}

af_err af_sort_by_key(af_array *out_keys, af_array *out_values,
                      const af_array keys, const af_array values,
                      const unsigned dim, const bool isAscending)
{
    try {
        ArrayInfo info = getInfo(keys);
        af_dtype type = info.getType();

        ArrayInfo vinfo = getInfo(values);

        DIM_ASSERT(3, info.elements() > 0);
        DIM_ASSERT(4, info.dims() == vinfo.dims());
        // Only Dim 0 supported
        ARG_ASSERT(5, dim == 0);

        af_array oKey;
        af_array oVal;

        switch(type) {
            case f32: sort_by_key_tmplt<float  >(&oKey, &oVal, keys, values, dim, isAscending);  break;
            case f64: sort_by_key_tmplt<double >(&oKey, &oVal, keys, values, dim, isAscending);  break;
            case s32: sort_by_key_tmplt<int    >(&oKey, &oVal, keys, values, dim, isAscending);  break;
            case u32: sort_by_key_tmplt<uint   >(&oKey, &oVal, keys, values, dim, isAscending);  break;
            case u8:  sort_by_key_tmplt<uchar  >(&oKey, &oVal, keys, values, dim, isAscending);  break;
            case b8:  sort_by_key_tmplt<char   >(&oKey, &oVal, keys, values, dim, isAscending);  break;
            default:  TYPE_ERROR(1, type);
        }
        std::swap(*out_keys , oKey);
        std::swap(*out_values , oVal);
    }
    CATCHALL;

    return AF_SUCCESS;
}

static af_err af_arith_real(af_array *out, const af_array lhs, const af_array rhs, const bool batchMode)
{
    try {

        ArrayInfo linfo = getInfo(lhs);
        ArrayInfo rinfo = getInfo(rhs);

        dim4 odims = getOutDims(linfo.dims(), rinfo.dims(), batchMode);

        const af_dtype otype = implicit(linfo.getType(), rinfo.getType());
        af_array res;
        switch (otype) {
        case f32: res = arithOp<float  , op>(lhs, rhs, odims); break;
        case f64: res = arithOp<double , op>(lhs, rhs, odims); break;
        case s32: res = arithOp<int    , op>(lhs, rhs, odims); break;
        case u32: res = arithOp<uint   , op>(lhs, rhs, odims); break;
        case u8 : res = arithOp<uchar  , op>(lhs, rhs, odims); break;
        case b8 : res = arithOp<char   , op>(lhs, rhs, odims); break;
        case s64: res = arithOp<intl   , op>(lhs, rhs, odims); break;
        case u64: res = arithOp<uintl  , op>(lhs, rhs, odims); break;
        case s16: res = arithOp<short  , op>(lhs, rhs, odims); break;
        case u16: res = arithOp<ushort , op>(lhs, rhs, odims); break;
        default: TYPE_ERROR(0, otype);
        }

        std::swap(*out, res);
    }
    CATCHALL;
    return AF_SUCCESS;
}

af_err convolve(af_array *out, af_array signal, af_array filter)
{
    try {
        ArrayInfo sInfo = getInfo(signal);
        ArrayInfo fInfo = getInfo(filter);

        af_dtype stype  = sInfo.getType();

        dim4 sdims = sInfo.dims();
        dim4 fdims = fInfo.dims();

        dim_type batchDim = baseDim+1;
        ARG_ASSERT(1, (sdims.ndims()<=batchDim));
        ARG_ASSERT(2, (fdims.ndims()<=batchDim));

        ConvolveBatchKind convBT = identifyBatchKind<baseDim>(sdims.ndims(), fdims.ndims());

        af_array output;
        switch(stype) {
            case c32: output = convolve<cfloat ,  cfloat, baseDim, expand>(signal, filter, convBT); break;
            case c64: output = convolve<cdouble, cdouble, baseDim, expand>(signal, filter, convBT); break;
            case f32: output = convolve<float  ,   float, baseDim, expand>(signal, filter, convBT); break;
            case f64: output = convolve<double ,  double, baseDim, expand>(signal, filter, convBT); break;
            case u32: output = convolve<uint   ,   float, baseDim, expand>(signal, filter, convBT); break;
            case s32: output = convolve<int    ,   float, baseDim, expand>(signal, filter, convBT); break;
            case u8:  output = convolve<uchar  ,   float, baseDim, expand>(signal, filter, convBT); break;
            case b8:  output = convolve<char   ,   float, baseDim, expand>(signal, filter, convBT); break;
            default: TYPE_ERROR(1, stype);
        }
        std::swap(*out,output);
    }
    CATCHALL;

    return AF_SUCCESS;
}

af_err af_replace_scalar(af_array a, const af_array cond, const double b)
{
    try {
        ArrayInfo ainfo = getInfo(a);
        ArrayInfo cinfo = getInfo(cond);

        ARG_ASSERT(1, cinfo.getType() == b8);
        DIM_ASSERT(1, cinfo.ndims() == ainfo.ndims());

        dim4 adims = ainfo.dims();
        dim4 cdims = cinfo.dims();

        for (int i = 0; i < 4; i++) {
            DIM_ASSERT(1, cdims[i] == adims[i]);
        }

        switch (ainfo.getType()) {
        case f32: replace_scalar<float  >(a, cond, b); break;
        case f64: replace_scalar<double >(a, cond, b); break;
        case c32: replace_scalar<cfloat >(a, cond, b); break;
        case c64: replace_scalar<cdouble>(a, cond, b); break;
        case s32: replace_scalar<int    >(a, cond, b); break;
        case u32: replace_scalar<uint   >(a, cond, b); break;
        case s64: replace_scalar<intl   >(a, cond, b); break;
        case u64: replace_scalar<uintl  >(a, cond, b); break;
        case u8:  replace_scalar<uchar  >(a, cond, b); break;
        case b8:  replace_scalar<char   >(a, cond, b); break;
        default:  TYPE_ERROR(2, ainfo.getType());
        }

    } CATCHALL;
    return AF_SUCCESS;
}

af_err af_corrcoef(double *realVal, double *imagVal, const af_array X, const af_array Y)
{
    try {
        ArrayInfo xInfo = getInfo(X);
        ArrayInfo yInfo = getInfo(Y);
        dim4 xDims      = xInfo.dims();
        dim4 yDims      = yInfo.dims();
        af_dtype xType  = xInfo.getType();
        af_dtype yType  = yInfo.getType();

        ARG_ASSERT(2, (xType==yType));
        ARG_ASSERT(2, (xDims.ndims()==yDims.ndims()));

        for (dim_t i=0; i<xDims.ndims(); ++i)
            ARG_ASSERT(2, (xDims[i]==yDims[i]));

        switch(xType) {
            case f64: *realVal = corrcoef<double, double>(X, Y); break;
            case f32: *realVal = corrcoef<float , float >(X, Y); break;
            case s32: *realVal = corrcoef<int   , float >(X, Y); break;
            case u32: *realVal = corrcoef<uint  , float >(X, Y); break;
            case s64: *realVal = corrcoef<intl  , double>(X, Y); break;
            case u64: *realVal = corrcoef<uintl , double>(X, Y); break;
            case s16: *realVal = corrcoef<short , float >(X, Y); break;
            case u16: *realVal = corrcoef<ushort, float >(X, Y); break;
            case  u8: *realVal = corrcoef<uchar , float >(X, Y); break;
            case  b8: *realVal = corrcoef<char  , float >(X, Y); break;
            default : TYPE_ERROR(1, xType);
        }
    }
    CATCHALL;
    return AF_SUCCESS;
}

af_err af_cplx(af_array *out, const af_array in)
{
    try {

        ArrayInfo info = getInfo(in);
        af_dtype type = info.getType();

        if (type == c32 || type == c64) {
            AF_ERROR("Inputs to cplx2 can not be of complex type", AF_ERR_ARG);
        }

        af_array tmp;
        AF_CHECK(af_constant(&tmp,
                             0, info.ndims(),
                             info.dims().get(),
                             type));

        af_array res;
        switch (type) {

        case f32: res = cplx<cfloat , float >(in, tmp, info.dims()); break;
        case f64: res = cplx<cdouble, double>(in, tmp, info.dims()); break;

        default: TYPE_ERROR(0, type);
        }

        AF_CHECK(af_release_array(tmp));

        std::swap(*out, res);
    }
    CATCHALL;
    return AF_SUCCESS;
}

static
void assign(af_array &out, const unsigned &ndims, const af_seq *index, const af_array &in)
{
    ArrayInfo iInfo = getInfo(in);
    ArrayInfo oInfo = getInfo(out);
    af_dtype iType  = iInfo.getType();

    dim4 const outDs = oInfo.dims();
    dim4 const iDims = iInfo.dims();

    ARG_ASSERT(0, (outDs.ndims()>=iDims.ndims()));
    ARG_ASSERT(1, (outDs.ndims()>=(int)ndims));

    AF_CHECK(af_eval(out));

    vector<af_seq> index_(index, index+ndims);
    dim4 const oStrides = af::toStride(index_, outDs);

    dim4 oDims = af::toDims(index_, outDs);
    dim4 oOffsets = af::toOffset(index_, outDs);

    Array<T> *dst = createRefArray<T>(getArray<T>(out), oDims, oOffsets, oStrides);

    for (int i = 0; i < 4; i++) {
        if (oDims[i] != iDims[i])
            AF_ERROR("Size mismatch between input and output", AF_ERR_SIZE);
    }

    bool noCaseExecuted = true;
    if (isComplex) {
        noCaseExecuted = false;
        switch(iType) {
            case c64: copy<cdouble, T>(*dst, getArray<cdouble>(in), scalar<T>(0), 1.0);  break;
            case c32: copy<cfloat , T>(*dst, getArray<cfloat >(in), scalar<T>(0), 1.0);  break;
            default : noCaseExecuted = true; break;
        }
    }

    static const T ZERO = scalar<T>(0);
    if(noCaseExecuted) {
        noCaseExecuted = false;
        switch(iType) {
            case f64: copy<double , T>(*dst, getArray<double>(in), ZERO, 1.0);  break;
            case f32: copy<float  , T>(*dst, getArray<float >(in), ZERO, 1.0);  break;
            case s32: copy<int    , T>(*dst, getArray<int   >(in), ZERO, 1.0);  break;
            case u32: copy<uint   , T>(*dst, getArray<uint  >(in), ZERO, 1.0);  break;
            case u8 : copy<uchar  , T>(*dst, getArray<uchar >(in), ZERO, 1.0);  break;
            case b8 : copy<char   , T>(*dst, getArray<char  >(in), ZERO, 1.0);  break;
            default : noCaseExecuted = true; break;
        }
    }

    if (noCaseExecuted)
        TYPE_ERROR(1, iType);

    delete dst;
}

af_err af_nearest_neighbour(af_array* idx, af_array* dist,
        const af_array query, const af_array train,
        const dim_t dist_dim, const uint n_dist,
        const af_match_type dist_type)
{
    try {
        ArrayInfo qInfo = getInfo(query);
        ArrayInfo tInfo = getInfo(train);
        af_dtype qType  = qInfo.getType();
        af_dtype tType  = tInfo.getType();
        af::dim4 qDims  = qInfo.dims();
        af::dim4 tDims  = tInfo.dims();

        uint train_samples = (dist_dim == 0) ? 1 : 0;

        DIM_ASSERT(2, qDims[dist_dim] == tDims[dist_dim]);
        DIM_ASSERT(2, qDims[2] == 1 && qDims[3] == 1);
        DIM_ASSERT(3, tDims[2] == 1 && tDims[3] == 1);
        DIM_ASSERT(4, (dist_dim == 0 || dist_dim == 1));
        DIM_ASSERT(5, n_dist > 0 && n_dist <= (uint)tDims[train_samples]);
        ARG_ASSERT(6, dist_type == AF_SAD || dist_type == AF_SSD || dist_type == AF_SHD);
        TYPE_ASSERT(qType == tType);

        // For Hamming, only u8, u32 and u64 allowed.
        af_array oIdx;
        af_array oDist;

        if(dist_type == AF_SHD) {
            TYPE_ASSERT(qType == u8 || qType == u32 || qType == u64);
            switch(qType) {
                case u8:  nearest_neighbour<uchar, uint>(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case u32: nearest_neighbour<uint , uint>(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case u64: nearest_neighbour<uintl, uint>(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                default : TYPE_ERROR(1, qType);
            }
        } else {
            switch(qType) {
                case f32: nearest_neighbour<float , float >(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case f64: nearest_neighbour<double, double>(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case s32: nearest_neighbour<int   , int   >(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case u32: nearest_neighbour<uint  , uint  >(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case s64: nearest_neighbour<intl  , intl  >(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case u64: nearest_neighbour<uintl , uintl >(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                case u8:  nearest_neighbour<uchar , uint  >(&oIdx, &oDist, query, train, dist_dim, n_dist, dist_type); break;
                default : TYPE_ERROR(1, qType);
            }
        }
        std::swap(*idx, oIdx);
        std::swap(*dist, oDist);
    }
    CATCHALL;

    return AF_SUCCESS;
}

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)
{
    using namespace detail;

    try {
        ArrayInfo lhsInfo = getInfo(lhs, false, true);
        ArrayInfo rhsInfo = getInfo(rhs, true, true);

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

        af_dtype lhs_type = lhsInfo.getType();
        af_dtype rhs_type = rhsInfo.getType();

        if (!(optLhs == AF_MAT_NONE ||
              optLhs == AF_MAT_TRANS ||
              optLhs == AF_MAT_CTRANS)) {
            AF_ERROR("Using this property is not yet supported in matmul", AF_ERR_NOT_SUPPORTED);
        }

        if (!(optRhs == AF_MAT_NONE ||
              optRhs == AF_MAT_TRANS ||
              optRhs == AF_MAT_CTRANS)) {
            AF_ERROR("Using this property is not yet supported in matmul", AF_ERR_NOT_SUPPORTED);
        }


        if (lhsInfo.ndims() > 2 ||
            rhsInfo.ndims() > 2) {
            AF_ERROR("matmul can not be used in batch mode", AF_ERR_BATCH);
        }

        TYPE_ASSERT(lhs_type == rhs_type);
        af_array output = 0;

        int aColDim = (optLhs == AF_MAT_NONE) ? 1 : 0;
        int bRowDim = (optRhs == AF_MAT_NONE) ? 0 : 1;

        DIM_ASSERT(1, lhsInfo.dims()[aColDim] == rhsInfo.dims()[bRowDim]);

        switch(lhs_type) {
            case f32: output = matmul<float  >(lhs, rhs, optLhs, optRhs);   break;
            case c32: output = matmul<cfloat >(lhs, rhs, optLhs, optRhs);   break;
            case f64: output = matmul<double >(lhs, rhs, optLhs, optRhs);   break;
            case c64: output = matmul<cdouble>(lhs, rhs, optLhs, optRhs);   break;
            default:  TYPE_ERROR(1, lhs_type);
        }
        std::swap(*out, output);
    }
    CATCHALL
    return AF_SUCCESS;
}

af_err af_get_dims(dim_t *d0, dim_t *d1, dim_t *d2, dim_t *d3,
                   const af_array in)
{
    try {
        ArrayInfo info = getInfo(in);
        *d0 = info.dims()[0];
        *d1 = info.dims()[1];
        *d2 = info.dims()[2];
        *d3 = info.dims()[3];
    }
    CATCHALL
    return AF_SUCCESS;
}

af_err af_get_dims(dim_t *d0, dim_t *d1, dim_t *d2, dim_t *d3,
                   const af_array in)
{
    try {
        // Do not check for device mismatch
        ArrayInfo info = getInfo(in, false, false);
        *d0 = info.dims()[0];
        *d1 = info.dims()[1];
        *d2 = info.dims()[2];
        *d3 = info.dims()[3];
    }
    CATCHALL
    return AF_SUCCESS;
}

af_err af_topk(af_array *values, af_array *indices, const af_array in,
               const int k, const int dim, const af_topk_function order) {
    try {
        af::topkFunction ord = (order == AF_TOPK_DEFAULT ? AF_TOPK_MAX : order);

        ArrayInfo inInfo = getInfo(in);

        ARG_ASSERT(2, (inInfo.ndims() > 0));

        if (inInfo.elements() == 1) {
            dim_t dims[1]   = {1};
            af_err errValue = af_constant(indices, 0, 1, dims, u32);
            return errValue == AF_SUCCESS ? af_retain_array(values, in)
                                          : errValue;
        }

        int rdim     = dim;
        auto &inDims = inInfo.dims();

        if (rdim == -1) {
            for (dim_t d = 0; d < 4; d++) {
                if (inDims[d] > 1) {
                    rdim = d;
                    break;
                }
            }
        }

        ARG_ASSERT(2, (inInfo.dims()[rdim] >= k));
        ARG_ASSERT(4, (k <= 256));  // TODO(umar): Remove this limitation

        if (rdim != 0)
            AF_ERROR("topk is supported along dimenion 0 only.",
                     AF_ERR_NOT_SUPPORTED);

        af_dtype type = inInfo.getType();

        switch (type) {
            // TODO(umar): FIX RETURN VALUES HERE
            case f32: topk<float>(values, indices, in, k, rdim, ord); break;
            case f64: topk<double>(values, indices, in, k, rdim, ord); break;
            case u32: topk<uint>(values, indices, in, k, rdim, ord); break;
            case s32: topk<int>(values, indices, in, k, rdim, ord); break;
            default: TYPE_ERROR(1, type);
        }
    }
    CATCHALL;

    return AF_SUCCESS;
}

af_err af_select(af_array *out, const af_array cond, const af_array a, const af_array b)
{
    try {
        ArrayInfo ainfo = getInfo(a);
        ArrayInfo binfo = getInfo(b);
        ArrayInfo cinfo = getInfo(cond);

        if(cinfo.ndims() == 0) {
            return af_retain_array(out, cond);
        }

        ARG_ASSERT(2, ainfo.getType() == binfo.getType());
        ARG_ASSERT(1, cinfo.getType() == b8);

        DIM_ASSERT(1, cinfo.ndims() == std::min(ainfo.ndims(), binfo.ndims()));

        dim4 adims = ainfo.dims();
        dim4 bdims = binfo.dims();
        dim4 cdims = cinfo.dims();
        dim4 odims(1, 1, 1, 1);

        for (int i = 0; i < 4; i++) {
            DIM_ASSERT(1, cdims[i] == std::min(adims[i], bdims[i]));
            DIM_ASSERT(2, adims[i] == bdims[i] || adims[i] == 1 || bdims[i] == 1);
            odims[i] = std::max(adims[i], bdims[i]);
        }

        af_array res;

        switch (ainfo.getType()) {
        case f32: res = select<float  >(cond, a, b, odims); break;
        case f64: res = select<double >(cond, a, b, odims); break;
        case c32: res = select<cfloat >(cond, a, b, odims); break;
        case c64: res = select<cdouble>(cond, a, b, odims); break;
        case s32: res = select<int    >(cond, a, b, odims); break;
        case u32: res = select<uint   >(cond, a, b, odims); break;
        case s64: res = select<intl   >(cond, a, b, odims); break;
        case u64: res = select<uintl  >(cond, a, b, odims); break;
        case s16: res = select<short  >(cond, a, b, odims); break;
        case u16: res = select<ushort >(cond, a, b, odims); break;
        case u8:  res = select<uchar  >(cond, a, b, odims); break;
        case b8:  res = select<char   >(cond, a, b, odims); break;
        default:  TYPE_ERROR(2, ainfo.getType());
        }

        std::swap(*out, res);
    } CATCHALL;
    return AF_SUCCESS;
}