#ifndef HEX_CPY_DMA_H
#define HEX_CPY_DMA_H
#include "dma-queue.h"
#include "hex-common.h"
#include "htp-tensor.h"
#include <stddef.h>
#include <stdint.h>
static inline void cpy_dma_sametype_reshape_contig(dma_queue * dma_q,
dma_addr_t dst,
dma_addr_t src0,
uint32_t total_bytes) {
if (total_bytes == 0) {
return;
}
const uint32_t max_chunk = DMA_SAFE_CHUNK_SIZE;
while (total_bytes > 0) {
const uint32_t chunk = MIN(total_bytes, max_chunk);
if (!dma_queue_push(dma_q, dma_make_data(dst, src0), chunk, chunk, chunk, 1)) {
dma_queue_flush(dma_q);
dma_queue_push(dma_q, dma_make_data(dst, src0), chunk, chunk, chunk, 1);
}
dst += chunk;
src0 += chunk;
total_bytes -= chunk;
}
}
static inline void cpy_dma_push_2d_chunked(dma_queue * dma_q,
dma_addr_t dst,
dma_addr_t src,
size_t dst_stride,
size_t src_stride,
size_t row_size,
uint32_t nrows) {
if (row_size == 0 || nrows == 0) {
return;
}
while (nrows > 0) {
const uint32_t cur_rows = MIN(nrows, DMA_MAX_NROWS);
if (!dma_queue_push(dma_q, dma_make_data(dst, src), dst_stride, src_stride, row_size, cur_rows)) {
dma_queue_flush(dma_q);
dma_queue_push(dma_q, dma_make_data(dst, src), dst_stride, src_stride, row_size, cur_rows);
}
dst += cur_rows * dst_stride;
src += cur_rows * src_stride;
nrows -= cur_rows;
}
}
static inline void cpy_dma_sametype_sameshape(dma_queue * dma_q,
const struct htp_tensor * dst,
const struct htp_tensor * src0,
uint32_t elem_size) {
const uint32_t ne00 = src0->ne[0];
const uint32_t ne01 = src0->ne[1];
const uint32_t ne02 = src0->ne[2];
const uint32_t ne03 = src0->ne[3];
if (ne00 == 0 || ne01 == 0 || ne02 == 0 || ne03 == 0) {
return;
}
const uint32_t nb01 = src0->nb[1];
const uint32_t nb02 = src0->nb[2];
const uint32_t nb03 = src0->nb[3];
const uint32_t nb1 = dst->nb[1];
const uint32_t nb2 = dst->nb[2];
const uint32_t nb3 = dst->nb[3];
const bool contiguous = htp_tensor_is_contiguous(src0, elem_size) && htp_tensor_is_contiguous(dst, elem_size);
if (contiguous) {
cpy_dma_sametype_reshape_contig(dma_q, dst->data, src0->data, ne00 * elem_size * ne01 * ne02 * ne03);
return;
}
uint32_t exp_src = ne01 * nb01;
uint32_t exp_dst = ne01 * nb1;
bool contiguous_outer = true;
if (ne02 != 1) {
contiguous_outer = contiguous_outer && (nb02 == exp_src) && (nb2 == exp_dst);
}
exp_src *= ne02;
exp_dst *= ne02;
if (ne03 != 1) {
contiguous_outer = contiguous_outer && (nb03 == exp_src) && (nb3 == exp_dst);
}
if (contiguous_outer) {
uint32_t total_rows = ne01 * ne02 * ne03;
cpy_dma_push_2d_chunked(dma_q, dst->data, src0->data, nb1, nb01, ne00 * elem_size, total_rows);
return;
}
for (uint32_t i03 = 0; i03 < ne03; i03++) {
for (uint32_t i02 = 0; i02 < ne02; i02++) {
dma_addr_t dst_data = dst->data + i02 * nb2 + i03 * nb3;
dma_addr_t src0_data = src0->data + i02 * nb02 + i03 * nb03;
cpy_dma_push_2d_chunked(dma_q, dst_data, src0_data, nb1, nb01, ne00 * elem_size, ne01);
}
}
}
#endif