Fix crash in xdna_aie_array initialization (amd#990)

Signed-off-by: Bikash Singha <bisingha@xcobisingha50x.amd.com>
Co-authored-by: Bikash Singha <bisingha@xcobisingha50x.amd.com>

Author: bisingha-xilinx

src/driver/amdxdna/ve2_debug.c

@@ -8,6 +8,7 @@
 #include <linux/completion.h>
 #include <linux/jiffies.h>
 #include <linux/sched.h>
+#include <linux/fdtable.h>
 
 #include "ve2_fw.h"
 #include "ve2_of.h"
@@ -675,6 +676,52 @@ static int ve2_get_array_async_error(struct amdxdna_dev *xdna, struct amdxdna_dr
 	return 0;
 }
 
+static int ve2_get_aie_part_fd(struct amdxdna_client *client,
+			       struct amdxdna_drm_get_array *args)
+{
+	struct amdxdna_dev *xdna = client->xdna;
+	struct amdxdna_ctx_priv *nhwctx;
+	struct amdxdna_ctx *ctx;
+	u32 hwctx_handle;
+	int srcu_idx;
+	int ret = 0;
+	int aie_fd;
+
+	hwctx_handle = args->num_element;
+	srcu_idx = srcu_read_lock(&client->ctx_srcu);
+	ctx = xa_load(&client->ctx_xa, hwctx_handle);
+	if (!ctx) {
+		XDNA_ERR(xdna, "Failed to get ctx %u", hwctx_handle);
+		ret = -EINVAL;
+		goto unlock;
+	}
+
+	nhwctx = ctx->priv;
+	if (!nhwctx || !nhwctx->aie_dev) {
+		XDNA_ERR(xdna, "AIE partition not available for hwctx %p", ctx);
+		ret = -ENODEV;
+		goto unlock;
+	}
+
+	aie_fd = aie_partition_get_fd(nhwctx->aie_dev);
+	if (aie_fd < 0) {
+		XDNA_ERR(xdna, "Failed to get AIE partition FD: %d", aie_fd);
+		ret = aie_fd;
+		goto unlock;
+	}
+
+	if (copy_to_user(u64_to_user_ptr(args->buffer), &aie_fd, sizeof(aie_fd))) {
+		XDNA_ERR(xdna, "Failed to copy AIE partition FD to user");
+		close_fd(aie_fd);
+		ret = -EFAULT;
+		goto unlock;
+	}
+
+unlock:
+	srcu_read_unlock(&client->ctx_srcu, srcu_idx);
+	return ret;
+}
+
 int ve2_get_array(struct amdxdna_client *client, struct amdxdna_drm_get_array *args)
 {
 	struct amdxdna_dev *xdna = client->xdna;
@@ -699,6 +746,9 @@ int ve2_get_array(struct amdxdna_client *client, struct amdxdna_drm_get_array *a
 	case DRM_AMDXDNA_HW_LAST_ASYNC_ERR:
 		ret = ve2_get_array_async_error(xdna, args);
 		break;
+	case DRM_AMDXDNA_HWCTX_AIE_PART_FD:
+		ret = ve2_get_aie_part_fd(client, args);
+		break;
 	default:
 		XDNA_ERR(xdna, "Not supported request parameter %u", args->param);
 		ret = -EOPNOTSUPP;

src/include/uapi/drm_local/amdxdna_accel.h

@@ -837,6 +837,7 @@ struct amdxdna_drm_get_array {
 #define DRM_AMDXDNA_FW_LOG_CONFIG	7
 #define DRM_AMDXDNA_FW_TRACE_CONFIG	8
 #define DRM_AMDXDNA_AIE_TILE_READ	9
+#define DRM_AMDXDNA_HWCTX_AIE_PART_FD		10
 	__u32 param; /* in */
 	__u32 element_size; /* in/out */
 #define AMDXDNA_MAX_NUM_ELEMENT			1024

src/shim_ve2/xdna_aie_array.cpp

@@ -49,6 +49,24 @@ get_driver_config(const pt::ptree& aie_meta)
   return driver_config;
 }
 
+int
+xdna_aie_array::
+get_aie_partition_fd(const xdna_hwctx* hwctx_obj)
+{
+  int aie_fd = -1;
+  auto dev = const_cast<xdna_hwctx*>(hwctx_obj)->get_device();
+
+  amdxdna_drm_get_array arg = {};
+  arg.param = DRM_AMDXDNA_HWCTX_AIE_PART_FD;
+  arg.element_size = sizeof(aie_fd);
+  arg.num_element = hwctx_obj->get_slotidx();  /* hwctx handle passed via num_element */
+  arg.buffer = reinterpret_cast<uintptr_t>(&aie_fd);
+
+  dev->get_edev()->ioctl(DRM_IOCTL_AMDXDNA_GET_ARRAY, &arg);
+
+  return aie_fd;
+}
+
 xdna_aie_array::
 xdna_aie_array(const xrt_core::device* device)
 {
@@ -72,7 +90,7 @@ xdna_aie_array(const xrt_core::device* device)
   int RC = XAie_GetPartitionFdList(&dev_inst_obj);
 
   if (RC != XAIE_OK) 
-   throw xrt_core::error(RC,"XAie_GetPartitionFdList failed \n");
+    throw xrt_core::error(RC, "XAie_GetPartitionFdList failed\n");
 
   XAie_List *NodePtr;
   XAie_PartitionList *ListNode;
@@ -83,10 +101,8 @@ xdna_aie_array(const xrt_core::device* device)
 
   int aie_part_fd = ListNode->PartitionFd;
 
-  //int aie_part_fd = fd;
-
   if (aie_part_fd < 0)
-    throw xrt_core::error(aie_part_fd,"fd is NEGATIVE\n");
+    throw xrt_core::error(aie_part_fd, "fd is NEGATIVE\n");
 
   fd = aie_part_fd;
   ConfigPtr.PartProp.Handle = fd;
@@ -125,22 +141,10 @@ xdna_aie_array(const xrt_core::device* device, const xdna_hwctx* hwctx_obj)
       throw xrt_core::error(-EINVAL, "Failed to setup AIE Partition: " + std::to_string(rc1));
   }
 
-  int RC = XAie_GetPartitionFdList(&dev_inst_obj);
-
-  if (RC != XAIE_OK) 
-   throw xrt_core::error(RC,"XAie_GetPartitionFdList failed \n");
-
-  XAie_List *NodePtr;
-  XAie_PartitionList *ListNode;
-
-  NodePtr = (XAie_List *)&dev_inst_obj.PartitionList.Next->Next;
-
-  ListNode = (XAie_PartitionList *)XAIE_CONTAINER_OF(NodePtr, XAie_PartitionList, Node);
-
-  int aie_part_fd = ListNode->PartitionFd;
-
+  // Get AIE partition FD from kernel via ioctl
+  int aie_part_fd = get_aie_partition_fd(hwctx_obj);
   if (aie_part_fd < 0)
-    throw xrt_core::error(aie_part_fd,"fd is NEGATIVE\n");
+    throw xrt_core::error(aie_part_fd, "Failed to get AIE partition FD\n");
 
   fd = aie_part_fd;
   ConfigPtr.PartProp.Handle = fd;

src/shim_ve2/xdna_aie_array.h

@@ -28,6 +28,7 @@ class xdna_aie_array {
 	XAie_DevInst *get_dev();
    adf::driver_config get_driver_config_hwctx(const xrt_core::device* device, const xdna_hwctx* hwctx);
 private:
+  int get_aie_partition_fd(const xdna_hwctx* hwctx_obj);
   int num_cols;
   int fd;
   XAie_DevInst* dev_inst;         // AIE Device Instance pointer

src/shim_ve2/xdna_hwctx.cpp

@@ -47,10 +47,11 @@ get_partition_info_main(const xrt_core::device* device,const pt::ptree& aie_meta
   info.base_address = aie_meta.get<uint64_t>("aie_metadata.driver_config.base_address");
 
   bool partinfo_found = false;
+  pid_t pid = getpid();
   auto data = xrt_core::device_query_default<xrt_core::query::aie_partition_info>(device, {});
 
   for (const auto& entry : data) {
-    if ( std::stoi(entry.metadata.id) == hw_context_id) {
+    if (entry.pid == pid && std::stoi(entry.metadata.id) == hw_context_id) {
       info.num_columns = entry.num_cols;
       info.start_column = entry.start_col;
       info.partition_id = (entry.num_cols << 8U) | (entry.start_col & 0xffU);
@@ -152,6 +153,11 @@ xdna_hwctx::
       return;
 
     m_hwq->unbind_hwctx();
+
+    // Explicitly destroy the aie_array before destroying the hw context
+    if(m_aie_array)
+      m_aie_array.reset();
+
     struct amdxdna_drm_destroy_hwctx arg = {};
     arg.handle = m_handle;
     m_device->get_edev()->ioctl(DRM_IOCTL_AMDXDNA_DESTROY_HWCTX, &arg);