dma.h

Go to the documentation of this file.

// Copyright 2023 ETH Zurich and University of Bologna.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0
 
#ifndef SNRT_SUPPORTS_DMA
#include <string.h>
#endif
 
#pragma once
 
#include <math.h>
 
#include "idma_compute.h"
 
inline uint32_t snrt_dma_start_1d(uint64_t dst, uint64_t src, size_t size,
                                  uint32_t channel) {
#ifdef SNRT_SUPPORTS_DMA
    uint32_t dst_lo = dst & 0xFFFFFFFF;
    uint32_t dst_hi = dst >> 32;
    uint32_t src_lo = src & 0xFFFFFFFF;
    uint32_t src_hi = src >> 32;
    uint32_t cfg = (channel << 2) | 0b00;
    uint32_t txid;
 
    asm volatile(
        "dmsrc %[src_lo], %[src_hi] \n"
        "dmdst %[dst_lo], %[dst_hi] \n"
        "dmcpy %[txid], %[size], %[cfg] \n"
        : [ txid ] "=r"(txid)
        : [ src_lo ] "r"(src_lo), [ src_hi ] "r"(src_hi),
          [ dst_lo ] "r"(dst_lo), [ dst_hi ] "r"(dst_hi), [ size ] "r"(size),
          [ cfg ] "r"(cfg));
    return txid;
#else
    memcpy((void *)dst, (const void *)src, size);
    return 0;
#endif
}
 
inline uint32_t snrt_dma_start_1d(volatile void *dst, volatile void *src,
                                  size_t size, uint32_t channel = 0) {
    return snrt_dma_start_1d((uint64_t)dst, (uint64_t)src, size, channel);
}
 
inline void snrt_dma_set_awuser(uint64_t field) {
#ifdef SNRT_SUPPORTS_DMA
    uint32_t user_low = (uint32_t)(field);
    uint32_t user_high = (uint32_t)(field >> 32);
    asm volatile("dmuser %[user_low], %[user_high] \n"
                 :
                 : [ user_low ] "r"(user_low), [ user_high ] "r"(user_high));
#endif
}
 
inline void snrt_dma_enable_multicast(uint64_t mask) {
    snrt_collective_t op;
    op.f.opcode = SNRT_COLLECTIVE_MULTICAST;
    op.f.mask = mask;
    snrt_dma_set_awuser(op.w);
}
 
inline void snrt_dma_enable_reduction(uint64_t mask,
                                      snrt_collective_opcode_t opcode) {
    snrt_collective_t op;
    op.f.opcode = opcode;
    op.f.mask = mask;
    snrt_dma_set_awuser(op.w);
}
 
inline void snrt_dma_disable_multicast() { snrt_dma_set_awuser(0); }
 
inline void snrt_dma_disable_reduction() { snrt_dma_set_awuser(0); }
 
inline uint32_t snrt_dma_start_1d_reduction(uint64_t dst, uint64_t src,
                                            size_t size, uint64_t mask,
                                            snrt_collective_opcode_t opcode,
                                            uint32_t channel = 0) {
    snrt_dma_enable_reduction(mask, opcode);
    uint32_t txid = snrt_dma_start_1d(dst, src, size, channel);
    snrt_dma_disable_reduction();
    return txid;
}
 
inline uint32_t snrt_dma_start_1d_reduction(uint64_t dst, uint64_t src,
                                            size_t size, snrt_comm_t comm,
                                            snrt_collective_opcode_t opcode,
                                            uint32_t channel = 0) {
    uint64_t mask = snrt_get_collective_mask(comm);
    uint32_t txid =
        snrt_dma_start_1d_reduction(dst, src, size, mask, opcode, channel);
    return txid;
}
 
inline uint32_t snrt_dma_start_1d_mcast(uint64_t dst, uint64_t src, size_t size,
                                        uint64_t mask, uint32_t channel = 0) {
    snrt_dma_enable_multicast(mask);
    uint32_t txid = snrt_dma_start_1d(dst, src, size, channel);
    snrt_dma_disable_multicast();
    return txid;
}
 
inline uint32_t snrt_dma_start_1d_mcast(uint64_t dst, uint64_t src, size_t size,
                                        snrt_comm_t comm,
                                        uint32_t channel = 0) {
    uint64_t mask = snrt_get_collective_mask(comm);
    uint32_t txid = snrt_dma_start_1d_mcast(dst, src, size, mask, channel);
    return txid;
}
 
inline uint32_t snrt_dma_start_1d_reduction(volatile void *dst,
                                            volatile void *src, size_t size,
                                            uint64_t mask,
                                            snrt_collective_opcode_t opcode,
                                            uint32_t channel = 0) {
    return snrt_dma_start_1d_reduction((uint64_t)dst, (uint64_t)src, size, mask,
                                       opcode, channel);
}
 
inline uint32_t snrt_dma_start_1d_mcast(volatile void *dst, volatile void *src,
                                        size_t size, uint64_t mask,
                                        uint32_t channel = 0) {
    return snrt_dma_start_1d_mcast((uint64_t)dst, (uint64_t)src, size, mask,
                                   channel);
}
 
inline snrt_dma_txid_t snrt_dma_start_2d(uint64_t dst, uint64_t src,
                                         size_t size, size_t dst_stride,
                                         size_t src_stride, size_t repeat,
                                         uint32_t channel) {
#ifdef SNRT_SUPPORTS_DMA
    uint32_t dst_lo = dst & 0xFFFFFFFF;
    uint32_t dst_hi = dst >> 32;
    uint32_t src_lo = src & 0xFFFFFFFF;
    uint32_t src_hi = src >> 32;
    uint32_t cfg = (channel << 2) | 0b10;
    uint32_t txid;
 
    asm volatile(
        "dmsrc %[src_lo], %[src_hi] \n"
        "dmdst %[dst_lo], %[dst_hi] \n"
        "dmstr %[src_stride], %[dst_stride] \n"
        "dmrep %[repeat] \n"
        "dmcpy %[txid], %[size], %[cfg] \n"
        : [ txid ] "=r"(txid)
        : [ src_lo ] "r"(src_lo), [ src_hi ] "r"(src_hi),
          [ dst_lo ] "r"(dst_lo), [ dst_hi ] "r"(dst_hi),
          [ dst_stride ] "r"(dst_stride), [ src_stride ] "r"(src_stride),
          [ repeat ] "r"(repeat), [ size ] "r"(size), [ cfg ] "r"(cfg));
 
    return txid;
#else
    // TODO(colluca): we can implement this as a series of memcpy calls
    return 0;
#endif
}
 
inline uint32_t snrt_dma_start_2d(volatile void *dst, volatile void *src,
                                  size_t size, size_t dst_stride,
                                  size_t src_stride, size_t repeat,
                                  uint32_t channel = 0) {
    return snrt_dma_start_2d((uint64_t)dst, (uint64_t)src, size, dst_stride,
                             src_stride, repeat, channel);
}
 
inline uint32_t snrt_dma_start_2d_mcast(uint64_t dst, uint64_t src, size_t size,
                                        size_t dst_stride, size_t src_stride,
                                        size_t repeat, uint32_t mask,
                                        uint32_t channel = 0) {
    snrt_dma_enable_multicast(mask);
    uint32_t txid = snrt_dma_start_2d(dst, src, size, dst_stride, src_stride,
                                      repeat, channel);
    snrt_dma_disable_multicast();
    return txid;
}
 
inline uint32_t snrt_dma_start_2d_mcast(volatile void *dst, volatile void *src,
                                        size_t size, size_t dst_stride,
                                        size_t src_stride, size_t repeat,
                                        uint32_t mask, uint32_t channel = 0) {
    return snrt_dma_start_2d_mcast((uint64_t)dst, (uint64_t)src, size,
                                   dst_stride, src_stride, repeat, mask,
                                   channel);
}
 
static inline uint32_t snrt_dma_busy(const uint32_t channel) {
#ifdef SNRT_SUPPORTS_DMA
    uint32_t busy;
    asm volatile("dmstati %[busy], (%[channel] << 2) | 2 \n"
                 : [ busy ] "=r"(busy)
                 : [ channel ] "i"(channel)
                 :);
    return busy;
#else
    return 0;
#endif
}
 
static inline uint32_t snrt_dma_would_block(const uint32_t channel) {
#ifdef SNRT_SUPPORTS_DMA
    uint32_t would_block;
    asm volatile("dmstati %[would_block], (%[channel] << 2) | 3 \n"
                 : [ would_block ] "=r"(would_block)
                 : [ channel ] "i"(channel)
                 :);
    return would_block;
#else
    return 0;
#endif
}
 
static inline void snrt_dma_wait(snrt_dma_txid_t txid,
                                 const uint32_t channel = 0) {
#ifdef SNRT_SUPPORTS_DMA
    asm volatile(
        "1: \n"
        "dmstati t0, (%[channel] << 2) | 0 \n"
        "bltu t0, %[txid], 1b \n"
        :
        : [ txid ] "r"(txid), [ channel ] "i"(channel)
        : "t0");
#endif
}
 
static inline void snrt_dma_wait_all(const uint32_t channel = 0) {
#ifdef SNRT_SUPPORTS_DMA
    while (snrt_dma_busy(channel))
        ;
#endif
}
 
inline void snrt_dma_wait_all_channels(uint32_t num_channels) {
    for (int c = 0; c < num_channels; c++) {
        snrt_dma_wait_all(c);
    }
}
 
inline void snrt_dma_start_tracking() {
#ifdef SNRT_SUPPORTS_DMA
    asm volatile("dmstati zero, 0 \n");
#endif
}
 
inline void snrt_dma_stop_tracking() {
#ifdef SNRT_SUPPORTS_DMA
    asm volatile("dmstati zero, 0 \n");
#endif
}
 
inline void snrt_dma_memset(void *ptr, uint8_t value, uint32_t len) {
#ifdef SNRT_SUPPORTS_DMA
    // We set the first 64 bytes to the value, and then we use the DMA to copy
    // these into the remaining memory region. DMA is used only if len is
    // larger than 64 bytes, and an integer multiple of 64 bytes.
    size_t n_1d_transfers = len / 64;
    size_t use_dma = (len % 64) == 0 && len > 64;
    uint8_t *p = (uint8_t *)ptr;
 
    uint32_t nbytes = len < 64 || !use_dma ? len : 64;
    while (nbytes--) {
        *p++ = value;
    }
 
    if (use_dma) {
        snrt_dma_start_2d(ptr, ptr, 64, 64, 0, n_1d_transfers);
        snrt_dma_wait_all();
    }
#else
    memset(ptr, (int)value, len);
#endif
}
 
inline snrt_dma_txid_t snrt_dma_load_1d_tile(volatile void *dst,
                                             volatile void *src,
                                             size_t tile_idx, size_t tile_size,
                                             uint32_t prec) {
    size_t tile_nbytes = tile_size * prec;
    return snrt_dma_start_1d(
        (uint64_t)dst, (uint64_t)src + tile_idx * tile_nbytes, tile_nbytes);
}
 
inline snrt_dma_txid_t snrt_dma_load_1d_tile_mcast(void *dst, void *src,
                                                   size_t tile_idx,
                                                   size_t tile_size,
                                                   uint32_t prec,
                                                   uint64_t mask) {
    size_t tile_nbytes = tile_size * prec;
    return snrt_dma_start_1d_mcast((uintptr_t)dst,
                                   (uintptr_t)src + tile_idx * tile_nbytes,
                                   tile_nbytes, mask);
}
 
inline snrt_dma_txid_t snrt_dma_reduction_load_1d_tile(
    void *dst, void *src, size_t tile_idx, size_t tile_size, uint32_t prec,
    uint64_t mask, snrt_collective_opcode_t opcode) {
    size_t tile_nbytes = tile_size * prec;
    return snrt_dma_start_1d_reduction((uintptr_t)dst,
                                       (uintptr_t)src + tile_idx * tile_nbytes,
                                       tile_nbytes, mask, opcode);
}
 
inline snrt_dma_txid_t snrt_dma_1d_to_2d(volatile void *dst, volatile void *src,
                                         size_t size, size_t row_size,
                                         size_t stride) {
    return snrt_dma_start_2d(dst, src, row_size, stride, row_size,
                             size / row_size);
}
 
inline snrt_dma_txid_t snrt_dma_2d_to_1d(volatile void *dst, volatile void *src,
                                         size_t size, size_t row_size,
                                         size_t stride) {
    return snrt_dma_start_2d(dst, src, row_size, row_size, stride,
                             size / row_size);
}
 
inline snrt_dma_txid_t snrt_dma_store_1d_tile(void *dst, void *src,
                                              size_t tile_idx, size_t tile_size,
                                              uint32_t prec) {
    size_t tile_nbytes = tile_size * prec;
    return snrt_dma_start_1d((uint64_t)dst + tile_idx * tile_nbytes,
                             (uint64_t)src, tile_nbytes);
}
 
inline snrt_dma_txid_t snrt_dma_load_2d_tile(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec, size_t tile_ld) {
    size_t src_offset = 0;
    // Advance src array in x0 and x1 dimensions, and convert to byte offset
    src_offset += tile_x0_idx * tile_x0_size;
    src_offset += tile_x1_idx * tile_x1_size * full_x0_size;
    src_offset *= prec;
    // Initiate transfer
    return snrt_dma_start_2d((uint64_t)dst,               // dst
                             (uint64_t)src + src_offset,  // src
                             tile_x0_size * prec,         // size
                             tile_ld,                     // dst_stride
                             full_x0_size * prec,         // src_stride
                             tile_x1_size                 // repeat
    );
}
 
inline snrt_dma_txid_t snrt_dma_load_2d_tile(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec) {
    return snrt_dma_load_2d_tile(dst, src, tile_x1_idx, tile_x0_idx,
                                 tile_x1_size, tile_x0_size, full_x0_size, prec,
                                 tile_x0_size * prec);
}
 
inline snrt_dma_txid_t snrt_dma_load_2d_tile_mcast(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec, size_t tile_ld, uint32_t mask) {
    size_t src_offset = 0;
    // Advance src array in x0 and x1 dimensions, and convert to byte offset
    src_offset += tile_x0_idx * tile_x0_size;
    src_offset += tile_x1_idx * tile_x1_size * full_x0_size;
    src_offset *= prec;
    // Initiate transfer
    return snrt_dma_start_2d_mcast((uint64_t)dst,               // dst
                                   (uint64_t)src + src_offset,  // src
                                   tile_x0_size * prec,         // size
                                   tile_ld,                     // dst_stride
                                   full_x0_size * prec,         // src_stride
                                   tile_x1_size,                // repeat
                                   mask                         // mask
    );
}
 
inline snrt_dma_txid_t snrt_dma_load_2d_tile_mcast(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec, uint32_t mask) {
    return snrt_dma_load_2d_tile_mcast(dst, src, tile_x1_idx, tile_x0_idx,
                                       tile_x1_size, tile_x0_size, full_x0_size,
                                       prec, tile_x0_size * prec, mask);
}
 
inline snrt_dma_txid_t snrt_dma_load_2d_tile_mcast(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec, snrt_comm_t comm) {
    uint64_t mask = snrt_get_collective_mask(comm);
    return snrt_dma_load_2d_tile_mcast(dst, src, tile_x1_idx, tile_x0_idx,
                                       tile_x1_size, tile_x0_size, full_x0_size,
                                       prec, mask);
}
 
inline snrt_dma_txid_t snrt_dma_load_2d_tile_in_banks(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec, size_t num_banks) {
    // Calculate new tile size after reshaping the tile in the selected banks
    size_t tile_x0_size_in_banks = (num_banks * SNRT_TCDM_BANK_WIDTH) / prec;
    size_t tile_x1_size_in_banks =
        ceil((tile_x1_size * tile_x0_size) / (double)tile_x0_size_in_banks);
    size_t tile_ld = SNRT_TCDM_HYPERBANK_WIDTH;
    return snrt_dma_load_2d_tile(dst, src, tile_x1_idx, tile_x0_idx,
                                 tile_x1_size_in_banks, tile_x0_size_in_banks,
                                 full_x0_size, prec, tile_ld);
}
 
inline snrt_dma_txid_t snrt_dma_store_2d_tile(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec, size_t tile_ld) {
    size_t dst_offset = 0;
    // Advance dst array in x0 and x1 dimensions, and convert to byte offset
    dst_offset += tile_x0_idx * tile_x0_size;
    dst_offset += tile_x1_idx * tile_x1_size * full_x0_size;
    dst_offset *= prec;
    // Initiate transfer
    return snrt_dma_start_2d((uint64_t)dst + dst_offset,  // dst
                             (uint64_t)src,               // src
                             tile_x0_size * prec,         // size
                             full_x0_size * prec,         // dst_stride
                             tile_ld,                     // src_stride
                             tile_x1_size                 // repeat
    );
}
 
inline snrt_dma_txid_t snrt_dma_store_2d_tile(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec) {
    return snrt_dma_store_2d_tile(dst, src, tile_x1_idx, tile_x0_idx,
                                  tile_x1_size, tile_x0_size, full_x0_size,
                                  prec, tile_x0_size * prec);
}
 
inline snrt_dma_txid_t snrt_dma_store_2d_tile_from_banks(
    void *dst, void *src, size_t tile_x1_idx, size_t tile_x0_idx,
    size_t tile_x1_size, size_t tile_x0_size, size_t full_x0_size,
    uint32_t prec, size_t num_banks) {
    // Calculate new tile size after reshaping the tile in the selected banks
    size_t tile_x0_size_in_banks = (num_banks * SNRT_TCDM_BANK_WIDTH) / prec;
    size_t tile_x1_size_in_banks =
        ceil((tile_x1_size * tile_x0_size) / (double)tile_x0_size_in_banks);
    size_t tile_ld = SNRT_TCDM_HYPERBANK_WIDTH;
    return snrt_dma_store_2d_tile(dst, src, tile_x1_idx, tile_x0_idx,
                                  tile_x1_size_in_banks, tile_x0_size_in_banks,
                                  full_x0_size, prec, tile_ld);
}
 
inline void snrt_dma_set_opcode_params(uint32_t opcode, uint32_t params) {
#ifdef SNRT_SUPPORTS_DMA_COMPUTE
    asm volatile("dmopc %[opcode], %[params] \n"
                 :
                 : [ opcode ] "r"(opcode), [ params ] "r"(params)
                 : "memory");
#endif
}
 
inline void snrt_dma_set_opcode(uint32_t opcode) {
    snrt_dma_set_opcode_params(opcode, 0u);
}
 
inline void snrt_dma_enable_transpose(uint32_t mode, uint32_t tensor_m,
                                      uint32_t tensor_n) {
    snrt_dma_set_opcode_params(
        IDMA_DMOPC_OPC_TRANSPOSE | ((mode & IDMA_DMOPC_RS1_TP_MODE_MASK)
                                    << IDMA_DMOPC_RS1_TP_MODE_SHIFT),
        ((tensor_m & IDMA_DMOPC_RS2_TP_TENSOR_M_MASK)
         << IDMA_DMOPC_RS2_TP_TENSOR_M_SHIFT) |
            ((tensor_n & IDMA_DMOPC_RS2_TP_TENSOR_N_MASK)
             << IDMA_DMOPC_RS2_TP_TENSOR_N_SHIFT));
}
 
inline void snrt_dma_disable_compute() {
    snrt_dma_set_opcode(IDMA_DMOPC_OPC_PASSTHROUGH);
}
 
inline uint32_t snrt_dma_start_transpose(uint64_t dst, uint64_t src,
                                         size_t size, uint32_t mode,
                                         uint32_t tensor_m, uint32_t tensor_n,
                                         uint32_t channel = 0) {
    snrt_dma_enable_transpose(mode, tensor_m, tensor_n);
    uint32_t txid = snrt_dma_start_1d(dst, src, size, channel);
    snrt_dma_disable_compute();
    return txid;
}
 
inline uint32_t snrt_dma_start_transpose(volatile void *dst, volatile void *src,
                                         size_t size, uint32_t mode,
                                         uint32_t tensor_m, uint32_t tensor_n,
                                         uint32_t channel = 0) {
    return snrt_dma_start_transpose((uint64_t)dst, (uint64_t)src, size, mode,
                                    tensor_m, tensor_n, channel);
}