Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion dkms.conf
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
PACKAGE_NAME=aws-neuronx
PACKAGE_VERSION=2.29.0.0
PACKAGE_VERSION=2.30.2.0
BUILT_MODULE_NAME[0]="neuron"
MAKE[0]="make -C ${kernel_source_dir} M=${dkms_tree}/${PACKAGE_NAME}/${PACKAGE_VERSION}/build"
CLEAN="make -C ${kernel_source_dir} M=${dkms_tree}/${PACKAGE_NAME}/${PACKAGE_VERSION}/build clean"
Expand Down
94 changes: 80 additions & 14 deletions neuron_cdev.c
Original file line number Diff line number Diff line change
Expand Up @@ -779,7 +779,12 @@ static int ncdev_get_dmabuf_fd_v2(void *param)
return 0;

err_close_fd:
#if LINUX_VERSION_CODE > KERNEL_VERSION(5, 10, 192)
close_fd(dmabuf_fd);
#else
__close_fd(current->files, dmabuf_fd);
#endif

return -EFAULT;
}

Expand Down Expand Up @@ -864,6 +869,8 @@ static int ncdev_mem_copy(struct neuron_device *nd, unsigned int cmd, void *para
u64 src_offset;
u64 dst_offset;
u64 size;
u64 sequence_num = 0;
void *context = NULL;
int ret;

if (cmd == NEURON_IOCTL_MEM_COPY) {
Expand All @@ -876,7 +883,16 @@ static int ncdev_mem_copy(struct neuron_device *nd, unsigned int cmd, void *para
src_offset = arg.src_offset;
dst_offset = arg.dst_offset;
size = arg.size;
dst_mc = ncdev_mem_handle_to_mem_chunk(nd, arg.dst_mem_handle);
} else if (cmd == NEURON_IOCTL_MEM_COPY64_DEPRECATED) {
struct neuron_ioctl_mem_copy64_deprecated arg;
ret = neuron_copy_from_user(__func__, &arg, (struct neuron_ioctl_mem_copy64_deprecated *)param, sizeof(arg));
if (ret)
return ret;
src_mem_handle = arg.src_mem_handle;
dst_mem_handle = arg.dst_mem_handle;
src_offset = arg.src_offset;
dst_offset = arg.dst_offset;
size = arg.size;
} else if (cmd == NEURON_IOCTL_MEM_COPY64) {
struct neuron_ioctl_mem_copy64 arg;
ret = neuron_copy_from_user(__func__, &arg, (struct neuron_ioctl_mem_copy64 *)param, sizeof(arg));
Expand All @@ -887,6 +903,8 @@ static int ncdev_mem_copy(struct neuron_device *nd, unsigned int cmd, void *para
src_offset = arg.src_offset;
dst_offset = arg.dst_offset;
size = arg.size;
sequence_num = arg.sequence_num;
context = arg.context;
} else {
return -EINVAL;
}
Expand All @@ -908,11 +926,21 @@ static int ncdev_mem_copy(struct neuron_device *nd, unsigned int cmd, void *para
}
ret = ndma_memcpy_mc(nd, src_mc, dst_mc, src_offset, dst_offset, size);
if (ret) {
pr_err("dma memcpy failed\n");
return ret;
}
if (sequence_num != 0) {
u32 nc_id = ndma_mc_pair_to_nc(src_mc, dst_mc);
int qid = ndhal->ndhal_ndmar.ndmar_get_h2t_def_qid(nc_id);
int submit_ret;

// Temporary fake async copy: the copy is already done, but report its result through CQE in queue order.
submit_ret = ndma_zerocopy_submit_completed(nd, nc_id, qid, sequence_num, 0, context);
if (submit_ret) {
return submit_ret;
}
}
trace_ioctl_mem_copy(nd, src_mc, dst_mc);
return 0;
return ret;
}

static int ncdev_mem_copy_async(struct neuron_device *nd, unsigned int cmd, void *param)
Expand Down Expand Up @@ -1256,6 +1284,8 @@ static int ncdev_mem_buf_zerocopy64(struct neuron_device *nd, unsigned int cmd,
u64 size;
u32 bar4_wr_threshold;
int h2t_qid;
u32 nc_id;
int qid;
int ret;
struct neuron_ioctl_mem_buf_copy64zc arg;
bool use_bar4_wr;
Expand Down Expand Up @@ -1292,6 +1322,13 @@ static int ncdev_mem_buf_zerocopy64(struct neuron_device *nd, unsigned int cmd,
return -EFAULT;
}

nc_id = ndma_mc_pair_to_nc(mc, mc);
qid = (h2t_qid == NEURON_DMA_H2T_DEFAULT_QID) ? ndhal->ndhal_ndmar.ndmar_get_h2t_def_qid(nc_id) : h2t_qid;
if (!ndmar_qid_valid(qid)) {
pr_err("nd%02d: invalid h2t queue index %d", nd->device_index, qid);
return -ENOENT;
}

// limit to internal threshold to prevent DoS attack
bar4_wr_threshold = (arg.bar4_wr_threshold < BAR4_WR_THRESHOLD_MAX) ? arg.bar4_wr_threshold : BAR4_WR_THRESHOLD_MAX;
use_bar4_wr = !narch_is_qemu() &&
Expand Down Expand Up @@ -1319,18 +1356,17 @@ static int ncdev_mem_buf_zerocopy64(struct neuron_device *nd, unsigned int cmd,
if (unlikely(ret)) {
ret = neuron_copy_from_user(__func__, nd->npdev.bar4 + cpy_offset, buffer, size);
}

// for async mode, BAR4 write is already done; enqueue a completed ctx to preserve CQE order.
// for sync mode, return the BAR4 write result directly
if (sequence_num != 0) {
ret = ndma_zerocopy_submit_completed(nd, nc_id, qid, sequence_num, ret, context);
}
return ret;
} else {
nrt_tensor_batch_op_t op;

u32 nc_id = ndma_mc_pair_to_nc(mc, mc);
int qid = (h2t_qid == NEURON_DMA_H2T_DEFAULT_QID) ? ndhal->ndhal_ndmar.ndmar_get_h2t_def_qid(nc_id) : h2t_qid;
dma_addr_t dev_base = ndma_mc_to_pa(mc); // the caller already does the range check for dev_base+offset

if (!ndmar_qid_valid(qid)) {
pr_err("nd%02d: invalid h2t queue index %d", nd->device_index, qid);
return -ENOENT;
}

if (!ndma_zerocopy_supported()) {
pr_err_once("nd%02d: zero copy is not supported for architectures requiring DMA retry", nd->device_index);
return -EINVAL;
Expand All @@ -1343,9 +1379,8 @@ static int ncdev_mem_buf_zerocopy64(struct neuron_device *nd, unsigned int cmd,
ret = ndma_zerocopy_submit(nd, nc_id, &op, 1, dev_base, qid,
copy_to_mem_handle ? true : false,
sequence_num, context);
return ret;
}

return ret;
}

static int ncdev_mem_buf_zerocopy64_batch(struct neuron_device *nd, void *param)
Expand Down Expand Up @@ -2027,7 +2062,8 @@ static long ncdev_driver_info(unsigned int cmd, void *param)
NEURON_DRIVER_FEATURE_ZEROCOPY | NEURON_DRIVER_FEATURE_PINNED_HOST_MEM |
NEURON_DRIVER_FEATURE_ALLOC_WITH_PA;
if (ndma_zerocopy_supported())
driver_info.feature_flags1 |= NEURON_DRIVER_FEATURE_ASYNC_IO;
driver_info.feature_flags1 |= NEURON_DRIVER_FEATURE_ASYNC_RW |
NEURON_DRIVER_FEATURE_ASYNC_COPY;

return copy_to_user(param, &driver_info, sizeof(driver_info));
}
Expand Down Expand Up @@ -3342,6 +3378,34 @@ inline static long ncdev_misc_ioctl(struct file *filep, unsigned int cmd, unsign
return -EINVAL;
}

static int ncdev_get_fw_bars(struct neuron_device *nd, void *param)
{
struct neuron_ioctl_get_bar_info arg;
int ret;
int i;

ret = neuron_copy_from_user(__func__, &arg, (struct neuron_ioctl_get_bar_info *)param, sizeof(arg));
if (ret) {
return ret;
}

if (arg.query_type == 0 || arg.query_type > NEURON_MAX_BAR_QUERY_TYPES) {
return -EINVAL;
}

if (!nd->bar_info[arg.query_type - 1].cached) {
return -EAGAIN;
}

arg.count = nd->bar_info[arg.query_type - 1].count;
for (i = 0; i < arg.count; i++) {
arg.bar_types[i] = nd->bar_info[arg.query_type - 1].bars[i].bar_type;
arg.bars[i] = nd->bar_info[arg.query_type - 1].bars[i].bar_address;
}

return copy_to_user(param, &arg, sizeof(arg));
}

static long ncdev_ioctl(struct file *filep, unsigned int cmd, unsigned long param)
{
struct ncdev *ncd;
Expand Down Expand Up @@ -3538,6 +3602,8 @@ static long ncdev_ioctl(struct file *filep, unsigned int cmd, unsigned long para
return ncdev_available_perf_profiles(nd, (void*)param);
} else if (cmd == NEURON_IOCTL_GET_ASYNC_H2T_DMA_COMPL_QUEUES) {
return ncdev_get_async_h2d_dma_compl_queues(nd, (void*)param);
} else if (cmd == NEURON_IOCTL_GET_BAR_INFO) {
return ncdev_get_fw_bars(nd, (void*)param);
}

// B/W compatibility
Expand Down
12 changes: 12 additions & 0 deletions neuron_device.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
#include "neuron_sysfs_metrics.h"
#include "neuron_log.h"
#include "neuron_power.h"
#include "neuron_fw_io.h"

#ifndef static_assert
#define static_assert(expr, ...)
Expand All @@ -42,6 +43,9 @@
// Global host memory buf size used for memset the device memory
#define MEMSET_HOST_BUF_SIZE MAX_DMA_DESC_SIZE // guessed optimal DMA transfer and PCIe TLP size.

// PCI BAR addresses on switched fabrics
#define NEURON_MAX_BAR_QUERY_TYPES 2

struct neuron_pci_device {
phys_addr_t bar0_pa;
void __iomem *bar0;
Expand All @@ -67,6 +71,12 @@ struct neuron_hbm_scrub_ctx {
struct mem_chunk *hostbuf_mc[NUM_DMA_ENG_PER_DEVICE];
};

struct neuron_device_bar_info {
bool cached;
u8 count;
struct fw_io_bar_entry bars[NEURON_MAX_SWITCH_BARS];
};

struct neuron_device {
struct pci_dev *pdev;
int device_index;
Expand Down Expand Up @@ -120,6 +130,8 @@ struct neuron_device {

struct neuron_hbm_scrub_ctx hbm_scrub_ctx;

struct neuron_device_bar_info bar_info[NEURON_MAX_BAR_QUERY_TYPES];

// volatile to prevent compiler optimizations since accessed by different threads
// Indicates whether any performance profile with 7200 Mhz HBM is supported by this device
volatile int supports_hbm_7200;
Expand Down
2 changes: 2 additions & 0 deletions neuron_dhal.h
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ struct ndhal_nc {
struct ndhal_nq {
u8 (*nnq_get_nqid) (struct neuron_device *nd, u8 nc_id, u8 index, u32 nq_type);
void (*nnq_set_hwaddr) (struct neuron_device *nd, u8 nc_id, u8 index, u32 nq_type, u32 size, u64 queue_pa);
u8 (*nnq_get_nq_queue_count) (void);
};

struct ndhal_mpset {
Expand Down Expand Up @@ -106,6 +107,7 @@ struct ndhal_fw_io {
int (*fw_io_read_csr_array) (void **addrs, u32 *values, u32 num_csrs, bool operational);
int (*fw_io_execute_request) (struct fw_io_ctx *ctx, u8 command_id, const u8 *req, u32 req_size, u8 *resp, u32 resp_size);
int (*fw_io_post_metric) (struct fw_io_ctx *ctx, u8 *data, u32 size);
void (*fw_io_cache_bar_info) (struct neuron_device *nd);
};

struct ndhal_mmap {
Expand Down
86 changes: 78 additions & 8 deletions neuron_dma.c
Original file line number Diff line number Diff line change
Expand Up @@ -949,7 +949,7 @@ enum ndma_zcdma_state {
NDMA_UNPINNED,
NDMA_PINNED_UNSUBMITTED,
NDMA_SUBMITTED,
NDMA_COMPLETED, // not in dma context queue anymore
NDMA_COMPLETED, // completed transfer waiting to be harvested from ctx queue to CQ
};

/* DMA context */
Expand Down Expand Up @@ -983,6 +983,8 @@ struct ndma_h2t_zcdma_context {

// Async-only
struct mm_struct *mm; // mm that owns the user buffers
s64 compl_ret; // result for in-order CQ completion of completed work. Only valid for async.
// used only for async BAR4 writes and temporary fake-async D2D copy
};

static void ndma_pinned_mem_process_release(struct kref *kref);
Expand All @@ -992,7 +994,7 @@ static void ndma_zc_release_ctx(struct ndma_h2t_zcdma_context *ctx, u64 *nr_pinn
// do not free or set completion_ptr null. it is managed by completion_pool in ctx queue
// do not free or set page_list null. it is managed by page_list_pool in ctx queue

if (ctx->state >= NDMA_PINNED_UNSUBMITTED) {
if (ctx->state >= NDMA_PINNED_UNSUBMITTED && ctx->nr_pages > 0) {
/* Only unpin if we pinned it ourselves (not pre-pinned memory) */
if (!ctx->prepin_proc) {
if (ctx->direction) {
Expand All @@ -1017,6 +1019,7 @@ static void ndma_zc_release_ctx(struct ndma_h2t_zcdma_context *ctx, u64 *nr_pinn
ctx->state = NDMA_INVALID;
ctx->sequence_num = 0;
ctx->context = NULL;
ctx->compl_ret = 0;
}

/* H2D DMA Completion Queue (CQ) */
Expand Down Expand Up @@ -1068,10 +1071,15 @@ void ndma_h2d_compl_queue_destroy(struct ndma_h2d_compl_queue *compl_queue)
compl_queue->tail = 0;
}

/*
* Emit the CQE for an accepted async io request. Caller must hold the ring
* lock that serializes compl_queue->tail. Submission failures return an error
* to userspace and must not call this function to emit a CQE.
*/
static void ndma_h2d_compl_queue_put(struct ndma_h2d_compl_queue *compl_queue,
u64 sequence_num,
s64 compl_ret,
void *context)
u64 sequence_num,
s64 compl_ret,
void *context)
{
u32 head = 0;
u32 tail = 0;
Expand Down Expand Up @@ -1178,7 +1186,7 @@ static void ndma_ctx_queue_inc_tail(struct ndma_ctx_queue *queue)
// Assume the ctx at old tail is already filled by caller
// Tail advance may also initialize/advance the pinned+unsubmitted and unpinned pointers
struct ndma_h2t_zcdma_context *ctx = &queue->entries[old_tail];
if (ctx->state == NDMA_PINNED_UNSUBMITTED) {
if (ctx->state == NDMA_PINNED_UNSUBMITTED || ctx->state == NDMA_COMPLETED) {
if (ndma_ctx_queue_pinned_unsubmitted_empty(queue)) {
// The first pinned+unsubmitted pointer appears at old_tail
queue->first_pinned_unsubmitted = old_tail;
Expand Down Expand Up @@ -1862,6 +1870,53 @@ int ndma_zerocopy_submit(struct neuron_device *nd,
return ret;
}

int ndma_zerocopy_submit_completed(struct neuron_device *nd, u32 nc_id, int qid,
u64 sequence_num, s64 compl_ret, void *context)
{
const int eng_id = ndhal->ndhal_ndmar.ndmar_get_h2t_eng_id(nd, nc_id);
struct ndma_eng *eng = &nd->ndma_engine[eng_id];
struct ndma_ring *ring = &eng->queues[qid].ring_info;
struct ndma_ctx_queue *ctx_queue = &ring->dma_ctx_queue;
struct ndma_h2t_zcdma_context *ctx;
int ret;

if (!ndmar_h2t_ring_is_owner(ring, nc_id)) {
pr_err("nd%02d: attempting to use qid %d that was not assigned to nc %d\n",
nd->device_index, qid, nc_id);
return -ENOENT;
}

ret = ndma_h2d_create_cmpltn_thread(nd);
if (ret) {
return ret;
}

mutex_lock(&ring->h2t_ring_lock);
if (ndma_ctx_queue_is_full(ctx_queue)) {
mutex_unlock(&ring->h2t_ring_lock);
return -EBUSY;
}

ctx = ndma_ctx_queue_peek_tail(ctx_queue);
ctx->eng = eng;
ctx->ring = ring;
ctx->direction = true;
ctx->last = true;
ctx->nr_pages = 0;
ctx->nr_desc = 0;
ctx->state = NDMA_COMPLETED;
ctx->sequence_num = sequence_num;
ctx->context = context;
ctx->compl_ret = compl_ret;
ndma_ctx_queue_inc_tail(ctx_queue);
mutex_unlock(&ring->h2t_ring_lock);

atomic64_or(BIT_ULL(ndhal->ndhal_ndmar.ndmar_ctx_queue_bit(eng_id, qid)),
&nd->dma_cmpltn_thread.nonempty_ctxq_bitmap);
wake_up(&nd->dma_cmpltn_thread.wait_queue);
return 0;
}

/* The completion flow for completion, remote pinning, and submission. Async IO only */
static int ndma_zerocopy_complete(struct neuron_device *nd,
struct ndma_eng *eng,
Expand Down Expand Up @@ -1897,7 +1952,13 @@ static int ndma_zerocopy_complete(struct neuron_device *nd,
}
struct ndma_h2t_zcdma_context *submitted_ctx = ndma_ctx_queue_pop_submitted(ctx_queue);

ret = ndma_memcpy_wait_for_completion(eng, ring, submitted_ctx->nr_desc + 1, submitted_ctx->completion_ptr, true, false);
if (submitted_ctx->state == NDMA_COMPLETED) {
// Dummy context: BAR4/fake-async copy completed before enqueue.
ret = submitted_ctx->compl_ret;
} else {
ret = ndma_memcpy_wait_for_completion(eng, ring, submitted_ctx->nr_desc + 1,
submitted_ctx->completion_ptr, true, false);
}
if (ret) {
err = ret;
pr_err("async h2d dma completion failed for seq num %llu: %d\n", submitted_ctx->sequence_num, ret);
Expand All @@ -1915,7 +1976,16 @@ static int ndma_zerocopy_complete(struct neuron_device *nd,
while (true) {
struct ndma_h2t_zcdma_context *pinned_unsubmitted_ctx = ndma_ctx_queue_peek_pinned_unsubmitted(ctx_queue);

if (!pinned_unsubmitted_ctx || !_ndma_zc_descs_available(eng, ring->qid, pinned_unsubmitted_ctx->nr_pages)) {
if (!pinned_unsubmitted_ctx) {
break;
}
if (pinned_unsubmitted_ctx->state == NDMA_COMPLETED) {
// Dummy context has no DMA descriptors; preserve CQ order only.
ndma_ctx_queue_inc_first_pinned_unsubmitted(ctx_queue);
did_work = true;
continue;
}
if (!_ndma_zc_descs_available(eng, ring->qid, pinned_unsubmitted_ctx->nr_pages)) {
break;
}

Expand Down
Loading