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

void transpose(Param<T> output, CParam<T> input) {
    const dim4 odims    = output.dims();
    const dim4 ostrides = output.strides();
    const dim4 istrides = input.strides();

    T *out            = output.get();
    T const *const in = input.get();

    for (dim_t l = 0; l < odims[3]; ++l) {
        for (dim_t k = 0; k < odims[2]; ++k) {
            // Outermost loop handles batch mode
            // if input has no data along third dimension
            // this loop runs only once
            for (dim_t j = 0; j < odims[1]; ++j) {
                for (dim_t i = 0; i < odims[0]; ++i) {
                    // calculate array indices based on offsets and strides
                    // the helper getIdx takes care of indices
                    const dim_t inIdx  = getIdx(istrides, j, i, k, l);
                    const dim_t outIdx = getIdx(ostrides, i, j, k, l);
                    if (conjugate)
                        out[outIdx] = getConjugate(in[inIdx]);
                    else
                        out[outIdx] = in[inIdx];
                }
            }
            // outData and inData pointers doesn't need to be
            // offset as the getIdx function is taking care
            // of the batch parameter
        }
    }
}

void transpose_inplace(Param<T> input) {
    const dim4 idims    = input.dims();
    const dim4 istrides = input.strides();

    T *in = input.get();

    for (dim_t l = 0; l < idims[3]; ++l) {
        for (dim_t k = 0; k < idims[2]; ++k) {
            // Outermost loop handles batch mode
            // if input has no data along third dimension
            // this loop runs only once
            //
            // Run only bottom triangle. std::swap swaps with upper triangle
            for (dim_t j = 0; j < idims[1]; ++j) {
                for (dim_t i = j + 1; i < idims[0]; ++i) {
                    // calculate array indices based on offsets and strides
                    // the helper getIdx takes care of indices
                    const dim_t iIdx = getIdx(istrides, j, i, k, l);
                    const dim_t oIdx = getIdx(istrides, i, j, k, l);
                    if (conjugate) {
                        in[iIdx] = getConjugate(in[iIdx]);
                        in[oIdx] = getConjugate(in[oIdx]);
                        std::swap(in[iIdx], in[oIdx]);
                    } else {
                        std::swap(in[iIdx], in[oIdx]);
                    }
                }
            }
        }
    }
}

void diff2(Param<T> out, CParam<T> in, int const dim)
{
    af::dim4 dims = out.dims();
    // Bool for dimension
    bool is_dim0 = dim == 0;
    bool is_dim1 = dim == 1;
    bool is_dim2 = dim == 2;
    bool is_dim3 = dim == 3;

    T const * const inPtr = in.get();
    T * outPtr = out.get();

    // TODO: Improve this
    for(dim_t l = 0; l < dims[3]; l++) {
        for(dim_t k = 0; k < dims[2]; k++) {
            for(dim_t j = 0; j < dims[1]; j++) {
                for(dim_t i = 0; i < dims[0]; i++) {
                    // Operation: out[index] = in[index + 1 * dim_size] - in[index]
                    int idx = getIdx(in.strides(), i, j, k, l);
                    int jdx = getIdx(in.strides(),
                            i + is_dim0, j + is_dim1,
                            k + is_dim2, l + is_dim3);
                    int kdx = getIdx(in.strides(),
                            i + 2 * is_dim0, j + 2 * is_dim1,
                            k + 2 * is_dim2, l + 2 * is_dim3);
                    int odx = getIdx(out.strides(), i, j, k, l);
                    outPtr[odx] = inPtr[kdx] + inPtr[idx] - inPtr[jdx] - inPtr[jdx];
                }
            }
        }
    }
}

void wrap_dim(Param<T> out, CParam<T> in, const dim_t wx, const dim_t wy,
              const dim_t sx, const dim_t sy, const dim_t px, const dim_t py) {
    const T* inPtr = in.get();
    T* outPtr      = out.get();

    af::dim4 idims    = in.dims();
    af::dim4 odims    = out.dims();
    af::dim4 istrides = in.strides();
    af::dim4 ostrides = out.strides();

    dim_t nx = (odims[0] + 2 * px - wx) / sx + 1;

    for (dim_t w = 0; w < idims[3]; w++) {
        for (dim_t z = 0; z < idims[2]; z++) {
            dim_t cIn      = w * istrides[3] + z * istrides[2];
            dim_t cOut     = w * ostrides[3] + z * ostrides[2];
            const T* iptr_ = inPtr + cIn;
            T* optr        = outPtr + cOut;

            for (dim_t col = 0; col < idims[d]; col++) {
                // Offset output ptr
                const T* iptr = iptr_ + col * istrides[d];

                // Calculate input window index
                dim_t winy = (col / nx);
                dim_t winx = (col % nx);

                dim_t startx = winx * sx;
                dim_t starty = winy * sy;

                dim_t spx = startx - px;
                dim_t spy = starty - py;

                // Short cut condition ensuring all values within input
                // dimensions
                bool cond = (spx >= 0 && spx + wx < odims[0] && spy >= 0 &&
                             spy + wy < odims[1]);

                for (dim_t y = 0; y < wy; y++) {
                    for (dim_t x = 0; x < wx; x++) {
                        dim_t xpad = spx + x;
                        dim_t ypad = spy + y;

                        dim_t iloc = (y * wx + x);
                        if (d == 0) iloc *= istrides[1];

                        if (cond || (xpad >= 0 && xpad < odims[0] &&
                                     ypad >= 0 && ypad < odims[1])) {
                            dim_t oloc =
                                (ypad * ostrides[1] + xpad * ostrides[0]);
                            // FIXME: When using threads, atomize this
                            optr[oloc] += iptr[iloc];
                        }
                    }
                }
            }
        }
    }
}

void assign(Param<T> out, af::dim4 dDims,
            CParam<T> rhs, std::vector<bool> const isSeq,
            std::vector<af_seq> const seqs, std::vector< CParam<uint> > idxArrs)
{
    af::dim4 pDims = out.dims();
    // retrieve dimensions & strides for array to which rhs is being copied to
    af::dim4 dst_offsets = toOffset(seqs, dDims);
    af::dim4 dst_strides = toStride(seqs, dDims);
    // retrieve rhs array dimenesions & strides
    af::dim4 src_dims    = rhs.dims();
    af::dim4 src_strides = rhs.strides();
    // declare pointers to af_array index data
    uint const * const ptr0 = idxArrs[0].get();
    uint const * const ptr1 = idxArrs[1].get();
    uint const * const ptr2 = idxArrs[2].get();
    uint const * const ptr3 = idxArrs[3].get();

    const T * src= rhs.get();
    T * dst      = out.get();

    for(dim_t l=0; l<src_dims[3]; ++l) {

        dim_t src_loff = l*src_strides[3];

        dim_t dst_lIdx = trimIndex(isSeq[3] ? l+dst_offsets[3] : ptr3[l], pDims[3]);
        dim_t dst_loff = dst_lIdx * dst_strides[3];

        for(dim_t k=0; k<src_dims[2]; ++k) {

            dim_t src_koff = k*src_strides[2];

            dim_t dst_kIdx = trimIndex(isSeq[2] ? k+dst_offsets[2] : ptr2[k], pDims[2]);
            dim_t dst_koff = dst_kIdx * dst_strides[2];

            for(dim_t j=0; j<src_dims[1]; ++j) {

                dim_t src_joff = j*src_strides[1];

                dim_t dst_jIdx = trimIndex(isSeq[1] ? j+dst_offsets[1] : ptr1[j], pDims[1]);
                dim_t dst_joff = dst_jIdx * dst_strides[1];

                for(dim_t i=0; i<src_dims[0]; ++i) {

                    dim_t src_ioff = i*src_strides[0];
                    dim_t src_idx  = src_ioff + src_joff + src_koff + src_loff;

                    dim_t dst_iIdx = trimIndex(isSeq[0] ? i+dst_offsets[0] : ptr0[i], pDims[0]);
                    dim_t dst_ioff = dst_iIdx * dst_strides[0];
                    dim_t dst_idx  = dst_ioff + dst_joff + dst_koff + dst_loff;

                    dst[dst_idx] = src[src_idx];
                }
            }
        }
    }
}

void select(Param<T> out, CParam<char> cond, CParam<T> a, CParam<T> b) {
    af::dim4 adims    = a.dims();
    af::dim4 astrides = a.strides();
    af::dim4 bdims    = b.dims();
    af::dim4 bstrides = b.strides();

    af::dim4 cdims    = cond.dims();
    af::dim4 cstrides = cond.strides();

    af::dim4 odims    = out.dims();
    af::dim4 ostrides = out.strides();

    bool is_a_same[] = {adims[0] == odims[0], adims[1] == odims[1],
                        adims[2] == odims[2], adims[3] == odims[3]};

    bool is_b_same[] = {bdims[0] == odims[0], bdims[1] == odims[1],
                        bdims[2] == odims[2], bdims[3] == odims[3]};

    bool is_c_same[] = {cdims[0] == odims[0], cdims[1] == odims[1],
                        cdims[2] == odims[2], cdims[3] == odims[3]};

    const T *aptr    = a.get();
    const T *bptr    = b.get();
    T *optr          = out.get();
    const char *cptr = cond.get();

    for (int l = 0; l < odims[3]; l++) {
        int o_off3 = ostrides[3] * l;
        int a_off3 = astrides[3] * is_a_same[3] * l;
        int b_off3 = bstrides[3] * is_b_same[3] * l;
        int c_off3 = cstrides[3] * is_c_same[3] * l;

        for (int k = 0; k < odims[2]; k++) {
            int o_off2 = ostrides[2] * k + o_off3;
            int a_off2 = astrides[2] * is_a_same[2] * k + a_off3;
            int b_off2 = bstrides[2] * is_b_same[2] * k + b_off3;
            int c_off2 = cstrides[2] * is_c_same[2] * k + c_off3;

            for (int j = 0; j < odims[1]; j++) {
                int o_off1 = ostrides[1] * j + o_off2;
                int a_off1 = astrides[1] * is_a_same[1] * j + a_off2;
                int b_off1 = bstrides[1] * is_b_same[1] * j + b_off2;
                int c_off1 = cstrides[1] * is_c_same[1] * j + c_off2;

                for (int i = 0; i < odims[0]; i++) {
                    bool cval = is_c_same[0] ? cptr[c_off1 + i] : cptr[c_off1];
                    T aval    = is_a_same[0] ? aptr[a_off1 + i] : aptr[a_off1];
                    T bval    = is_b_same[0] ? bptr[b_off1 + i] : bptr[b_off1];
                    T oval    = cval ? aval : bval;
                    optr[o_off1 + i] = oval;
                }
            }
        }
    }
}

void approx2(Param<InT> output, CParam<InT> input,
             CParam<LocT> xposition, CParam<LocT> yposition,
             float const offGrid, af_interp_type method)
{
    InT * out = output.get();
    const LocT *xpos = xposition.get();
    const LocT *ypos = yposition.get();

    af::dim4 const odims     = output.dims();
    af::dim4 const idims     = input.dims();
    af::dim4 const xdims     = xposition.dims();
    af::dim4 const ostrides  = output.strides();
    af::dim4 const istrides  = input.strides();
    af::dim4 const xstrides  = xposition.strides();
    af::dim4 const ystrides  = yposition.strides();

    Interp2<InT, LocT, order> interp;
    bool batch = !(xdims[2] == 1 && xdims[3] == 1);

    for(dim_t idw = 0; idw < odims[3]; idw++) {
        for(dim_t idz = 0; idz < odims[2]; idz++) {

            dim_t xoffzw = idw * xstrides[3] + idz * xstrides[2];
            dim_t yoffzw = idw * ystrides[3] + idz * ystrides[2];
            dim_t ooffzw = idw * ostrides[3] + idz * ostrides[2];
            dim_t ioffzw = idw * istrides[3] + idz * istrides[2];

            for(dim_t idy = 0; idy < odims[1]; idy++) {
                dim_t xoff = xoffzw * batch + idy * xstrides[1];
                dim_t yoff = yoffzw * batch + idy * ystrides[1];
                dim_t ooff = ooffzw         + idy * ostrides[1];

                for(dim_t idx = 0; idx < odims[0]; idx++) {

                    const LocT x = xpos[xoff + idx];
                    const LocT y = ypos[yoff + idx];

                    // FIXME: Only cubic interpolation is doing clamping
                    // We need to make it consistent across all methods
                    // Not changing the behavior because tests will fail
                    bool clamp = order == 3;

                    if (x < 0 || idims[0] < x + 1 ||
                        y < 0 || idims[1] < y + 1 ) {
                        out[ooff + idx] = scalar<InT>(offGrid);
                    } else {
                        interp(output, ooff + idx, input, ioffzw, x, y, method, 1, clamp);
                    }
                }
            }
        }
    }
}

void copyElemwise(Param<OutT> dst, CParam<InT> src, OutT default_value, double factor)
{
    af::dim4 src_dims       = src.dims();
    af::dim4 dst_dims       = dst.dims();
    af::dim4 src_strides    = src.strides();
    af::dim4 dst_strides    = dst.strides();

    InT const * const src_ptr = src.get();
    OutT * dst_ptr      = dst.get();

    dim_t trgt_l = std::min(dst_dims[3], src_dims[3]);
    dim_t trgt_k = std::min(dst_dims[2], src_dims[2]);
    dim_t trgt_j = std::min(dst_dims[1], src_dims[1]);
    dim_t trgt_i = std::min(dst_dims[0], src_dims[0]);

    for(dim_t l=0; l<dst_dims[3]; ++l) {

        dim_t src_loff = l*src_strides[3];
        dim_t dst_loff = l*dst_strides[3];
        bool isLvalid = l<trgt_l;

        for(dim_t k=0; k<dst_dims[2]; ++k) {

            dim_t src_koff = k*src_strides[2];
            dim_t dst_koff = k*dst_strides[2];
            bool isKvalid = k<trgt_k;

            for(dim_t j=0; j<dst_dims[1]; ++j) {

                dim_t src_joff = j*src_strides[1];
                dim_t dst_joff = j*dst_strides[1];
                bool isJvalid = j<trgt_j;

                for(dim_t i=0; i<dst_dims[0]; ++i) {
                    OutT temp = default_value;
                    if (isLvalid && isKvalid && isJvalid && i<trgt_i) {
                        dim_t src_idx = i*src_strides[0] + src_joff + src_koff + src_loff;
                        temp = OutT(src_ptr[src_idx])*OutT(factor);
                    }
                    dim_t dst_idx = i*dst_strides[0] + dst_joff + dst_koff + dst_loff;
                    dst_ptr[dst_idx] = temp;
                }
            }
        }
    }
}

void join(const int dim, Param<T> out, const std::vector<CParam<T>> inputs)
{
    af::dim4 zero(0,0,0,0);
    af::dim4 d = zero;
    switch(dim) {
        case 0:
            join_append<T, T, 0>(out.get(), inputs[0].get(), zero,
                        out.dims(), inputs[0].dims(), out.strides(), inputs[0].strides());
            for(int i = 1; i < n_arrays; i++) {
                d += inputs[i - 1].dims();
                join_append<T, T, 0>(out.get(), inputs[i].get(), calcOffset<0>(d),
                        out.dims(), inputs[i].dims(), out.strides(), inputs[i].strides());
            }
            break;
        case 1:
            join_append<T, T, 1>(out.get(), inputs[0].get(), zero,
                        out.dims(), inputs[0].dims(), out.strides(), inputs[0].strides());
            for(int i = 1; i < n_arrays; i++) {
                d += inputs[i - 1].dims();
                join_append<T, T, 1>(out.get(), inputs[i].get(), calcOffset<1>(d),
                        out.dims(), inputs[i].dims(), out.strides(), inputs[i].strides());
            }
            break;
        case 2:
            join_append<T, T, 2>(out.get(), inputs[0].get(), zero,
                        out.dims(), inputs[0].dims(), out.strides(), inputs[0].strides());
            for(int i = 1; i < n_arrays; i++) {
                d += inputs[i - 1].dims();
                join_append<T, T, 2>(out.get(), inputs[i].get(), calcOffset<2>(d),
                        out.dims(), inputs[i].dims(), out.strides(), inputs[i].strides());
            }
            break;
        case 3:
            join_append<T, T, 3>(out.get(), inputs[0].get(), zero,
                        out.dims(), inputs[0].dims(), out.strides(), inputs[0].strides());
            for(int i = 1; i < n_arrays; i++) {
                d += inputs[i - 1].dims();
                join_append<T, T, 3>(out.get(), inputs[i].get(), calcOffset<3>(d),
                        out.dims(), inputs[i].dims(), out.strides(), inputs[i].strides());
            }
            break;
    }
}

void select_scalar(Param<T> out, CParam<char> cond, CParam<T> a,
                   const double b) {
    af::dim4 astrides = a.strides();
    af::dim4 adims    = a.dims();
    af::dim4 cstrides = cond.strides();
    af::dim4 cdims    = cond.dims();

    af::dim4 odims    = out.dims();
    af::dim4 ostrides = out.strides();

    const T *aptr    = a.get();
    T *optr          = out.get();
    const char *cptr = cond.get();

    bool is_a_same[] = {adims[0] == odims[0], adims[1] == odims[1],
                        adims[2] == odims[2], adims[3] == odims[3]};

    bool is_c_same[] = {cdims[0] == odims[0], cdims[1] == odims[1],
                        cdims[2] == odims[2], cdims[3] == odims[3]};

    for (int l = 0; l < odims[3]; l++) {
        int o_off3 = ostrides[3] * l;
        int a_off3 = astrides[3] * is_a_same[3] * l;
        int c_off3 = cstrides[3] * is_c_same[3] * l;

        for (int k = 0; k < odims[2]; k++) {
            int o_off2 = ostrides[2] * k + o_off3;
            int a_off2 = astrides[2] * is_a_same[2] * k + a_off3;
            int c_off2 = cstrides[2] * is_c_same[2] * k + c_off3;

            for (int j = 0; j < odims[1]; j++) {
                int o_off1 = ostrides[1] * j + o_off2;
                int a_off1 = astrides[1] * is_a_same[1] * j + a_off2;
                int c_off1 = cstrides[1] * is_c_same[1] * j + c_off2;

                for (int i = 0; i < odims[0]; i++) {
                    bool cval = is_c_same[0] ? cptr[c_off1 + i] : cptr[c_off1];
                    T aval    = is_a_same[0] ? aptr[a_off1 + i] : aptr[a_off1];
                    optr[o_off1 + i] = (flip ^ cval) ? aval : b;
                }
            }
        }
    }
}

void lookup(Param<InT> out, CParam<InT> input, CParam<IndexT> indices,
            unsigned const dim) {
    const af::dim4 iDims    = input.dims();
    const af::dim4 oDims    = out.dims();
    const af::dim4 iStrides = input.strides();
    const af::dim4 oStrides = out.strides();
    const InT *inPtr        = input.get();
    const IndexT *idxPtr    = indices.get();

    InT *outPtr = out.get();

    for (dim_t l = 0; l < oDims[3]; ++l) {
        dim_t iLOff = iStrides[3] *
                      (dim == 3 ? trimIndex((dim_t)idxPtr[l], iDims[3]) : l);
        dim_t oLOff = l * oStrides[3];

        for (dim_t k = 0; k < oDims[2]; ++k) {
            dim_t iKOff =
                iStrides[2] *
                (dim == 2 ? trimIndex((dim_t)idxPtr[k], iDims[2]) : k);
            dim_t oKOff = k * oStrides[2];

            for (dim_t j = 0; j < oDims[1]; ++j) {
                dim_t iJOff =
                    iStrides[1] *
                    (dim == 1 ? trimIndex((dim_t)idxPtr[j], iDims[1]) : j);
                dim_t oJOff = j * oStrides[1];

                for (dim_t i = 0; i < oDims[0]; ++i) {
                    dim_t iIOff =
                        iStrides[0] *
                        (dim == 0 ? trimIndex((dim_t)idxPtr[i], iDims[0]) : i);
                    dim_t oIOff = i * oStrides[0];

                    outPtr[oLOff + oKOff + oJOff + oIOff] =
                        inPtr[iLOff + iKOff + iJOff + iIOff];
                }
            }
        }
    }
}

void diagExtract(Param<T> out, CParam<T> in, int const num)
{
    af::dim4 const odims = out.dims();
    af::dim4 const idims = in.dims();

    int const i_off = (num > 0) ? (num * in.strides(1)) : (-num);

    for (int l = 0; l < (int)odims[3]; l++) {

        for (int k = 0; k < (int)odims[2]; k++) {
            const T *iptr = in.get() + l * in.strides(3) + k * in.strides(2) + i_off;
            T *optr = out.get() + l * out.strides(3) + k * out.strides(2);

            for (int i = 0; i < (int)odims[0]; i++) {
                T val = scalar<T>(0);
                if (i < idims[0] && i < idims[1]) val =  iptr[i * in.strides(1) + i];
                optr[i] = val;
            }
        }
    }
}

void diagCreate(Param<T> out, CParam<T> in, int const num)
{
    int batch = in.dims(1);
    int size  = out.dims(0);

    T const * iptr = in.get();
    T * optr = out.get();

    for (int k = 0; k < batch; k++) {
        for (int j = 0; j < size; j++) {
            for (int i = 0; i < size; i++) {
                T val = scalar<T>(0);
                if (i == j - num) {
                    val = (num > 0) ? iptr[i] : iptr[j];
                }
                optr[i + j * out.strides(1)] = val;
            }
        }
        optr += out.strides(2);
        iptr += in.strides(1);
    }
}

    static void copy(Param<T> dst, CParam<T> src)
    {
        af::dim4 src_dims       = src.dims();
        af::dim4 dst_dims       = dst.dims();
        af::dim4 src_strides    = src.strides();
        af::dim4 dst_strides    = dst.strides();

        T const * src_ptr = src.get();
        T * dst_ptr       = dst.get();

        // find the major-most dimension, which is linear in both arrays
        int linear_end = 0;
        dim_t count = 1;
        while (linear_end < 4
                && count == src_strides[linear_end]
                && count == dst_strides[linear_end]) {
            count *= src_dims[linear_end];
            ++linear_end;
        }

        // traverse through the array using strides only until neccessary
        copy_go(dst_ptr, dst_strides, dst_dims, src_ptr, src_strides, src_dims, 3, linear_end);
    }

void fft_inplace(Param<T> in, const af::dim4 iDataDims)
{
    int t_dims[rank];
    int in_embed[rank];

    const af::dim4 idims = in.dims();

    computeDims<rank>(t_dims  , idims);
    computeDims<rank>(in_embed , iDataDims);

    const af::dim4 istrides = in.strides();

    typedef typename fftw_transform<T>::ctype_t ctype_t;
    typename fftw_transform<T>::plan_t plan;

    fftw_transform<T> transform;

    int batch = 1;
    for (int i = rank; i < 4; i++) {
        batch *= idims[i];
    }

    plan = transform.create(rank,
                            t_dims,
                            (int)batch,
                            (ctype_t *)in.get(),
                            in_embed, (int)istrides[0],
                            (int)istrides[rank],
                            (ctype_t *)in.get(),
                            in_embed, (int)istrides[0],
                            (int)istrides[rank],
                            direction ? FFTW_FORWARD : FFTW_BACKWARD,
                            FFTW_ESTIMATE);

    transform.execute(plan);
    transform.destroy(plan);
}