Method exec

dnn/src/aarch64/relayout/opr_impl.cpp:465–500

}  // namespace megdnn

void aarch64::RelayoutForwardImpl::exec(
        _megdnn_tensor_in src0, _megdnn_tensor_out dst0, Handle* src_handle) {
    check_cpu_handle(src_handle);
    TensorND src = src0, dst = dst0;
    check_layout_and_canonize(src.layout, dst.layout);

    // FIXME: optimize for lowbit cases
    if (src.layout.dtype.enumv() == DTypeEnum::QuantizedS4 ||
        src.layout.dtype.enumv() == DTypeEnum::Quantized4Asymm) {
        fallback::RelayoutForwardImpl::exec(src0, dst0, src_handle);
        return;
    }
    relayout::TransposeParam trans_param;
    bool trans = relayout::is_transpose(src.layout, dst.layout, trans_param, true);
    if (trans && trans_param.c == 1 && src0.layout.dtype.size() == 1) {
        MEGDNN_DISPATCH_CPU_KERN_OPR(transpose_fallback::transpose<TransposeByte>(
                trans_param.batch, trans_param.m, trans_param.n,
                static_cast<TransposeByte*>(src.raw_ptr()),
                static_cast<TransposeByte*>(dst.raw_ptr()), trans_param.stride_m));
        return;
    } else if (trans && trans_param.c == 1 && src0.layout.dtype.size() == 2) {
        MEGDNN_DISPATCH_CPU_KERN_OPR(transpose_fallback::transpose<Transpose2Byte>(
                trans_param.batch, trans_param.m, trans_param.n,
                static_cast<Transpose2Byte*>(src.raw_ptr()),
                static_cast<Transpose2Byte*>(dst.raw_ptr()), trans_param.stride_m));
        return;
    } else if (trans && trans_param.c == 1 && src0.layout.dtype.size() == 4) {
        MEGDNN_DISPATCH_CPU_KERN_OPR(transpose_fallback::transpose<Transpose4Byte>(
                trans_param.batch, trans_param.m, trans_param.n,
                static_cast<Transpose4Byte*>(src.raw_ptr()),
                static_cast<Transpose4Byte*>(dst.raw_ptr()), trans_param.stride_m));
        return;
    }

    exec_after_preprocess(src, dst, trans ? &trans_param : nullptr);
}

// vim: syntax=cpp.doxygen