[PATCH mlx5-next 10/10] mlx5: Move data direct implementation to mlx5_core

[Tariq Toukan <tariqt@nvidia.com>]

From: Dragos Tatulea <dtatulea@nvidia.com>

As data direct will also be used by netdev devices, move the data direct
code to mlx5_core. Now mlx5_core is responsible for data direct
initialization and resource creation.

Now data_direct is initialized in mlx5_load() before attaching consumers
and clean it up in mlx5_unload() after detaching them. The data direct
PCI driver is moved and registered in mlx5_core.

IB specific work stays in mlx5_ib: notifier registration, mutex
initialization and MR list initialization.

This patch has no functional changes.

Signed-off-by: Dragos Tatulea <dtatulea@nvidia.com>
---
 drivers/infiniband/hw/mlx5/Makefile           |  1 -
 drivers/infiniband/hw/mlx5/data_direct.h      | 29 ------
 drivers/infiniband/hw/mlx5/main.c             | 23 +----
 drivers/infiniband/hw/mlx5/mlx5_ib.h          |  3 +-
 drivers/infiniband/hw/mlx5/mr.c               |  7 +-
 drivers/infiniband/hw/mlx5/odp.c              |  2 +-
 drivers/infiniband/hw/mlx5/std_types.c        |  3 +-
 drivers/infiniband/hw/mlx5/umr.c              |  4 +-
 .../net/ethernet/mellanox/mlx5/core/Makefile  |  2 +-
 .../mellanox/mlx5/core}/data_direct.c         | 91 ++++++++++---------
 .../net/ethernet/mellanox/mlx5/core/main.c    | 15 +++
 .../ethernet/mellanox/mlx5/core/mlx5_core.h   |  5 +
 include/linux/mlx5/data_direct.h              | 23 +++--
 include/linux/mlx5/driver.h                   |  2 +
 14 files changed, 97 insertions(+), 113 deletions(-)
 delete mode 100644 drivers/infiniband/hw/mlx5/data_direct.h
 rename drivers/{infiniband/hw/mlx5 => net/ethernet/mellanox/mlx5/core}/data_direct.c (79%)

diff --git a/drivers/infiniband/hw/mlx5/Makefile b/drivers/infiniband/hw/mlx5/Makefile
index dd7bb377f491..35afc7bc57d2 100644
--- a/drivers/infiniband/hw/mlx5/Makefile
+++ b/drivers/infiniband/hw/mlx5/Makefile
@@ -6,7 +6,6 @@ mlx5_ib-y := ah.o \
 	     cong.o \
 	     counters.o \
 	     cq.o \
-	     data_direct.o \
 	     dm.o \
 	     dmah.o \
 	     doorbell.o \
diff --git a/drivers/infiniband/hw/mlx5/data_direct.h b/drivers/infiniband/hw/mlx5/data_direct.h
deleted file mode 100644
index 49cdd407018a..000000000000
--- a/drivers/infiniband/hw/mlx5/data_direct.h
+++ /dev/null
@@ -1,29 +0,0 @@
-/* SPDX-License-Identifier: GPL-2.0 OR Linux-OpenIB */
-/*
- * Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved
- */
-
-#ifndef _MLX5_IB_DATA_DIRECT_H
-#define _MLX5_IB_DATA_DIRECT_H
-
-#include <linux/notifier.h>
-#include <linux/mlx5/data_direct.h>
-
-struct mlx5_core_dev;
-struct mlx5_ib_dev;
-
-enum mlx5_data_direct_event {
-	MLX5_DATA_DIRECT_UNBIND,
-};
-
-struct mlx5_data_direct_dev {
-	struct device *device;
-	struct pci_dev *pdev;
-	char *vuid;
-	struct list_head list;
-};
-
-int mlx5_data_direct_driver_register(void);
-void mlx5_data_direct_driver_unregister(void);
-
-#endif
diff --git a/drivers/infiniband/hw/mlx5/main.c b/drivers/infiniband/hw/mlx5/main.c
index 941928a782bd..dd2a3996b14e 100644
--- a/drivers/infiniband/hw/mlx5/main.c
+++ b/drivers/infiniband/hw/mlx5/main.c
@@ -52,7 +52,6 @@
 #include <rdma/mlx5_user_ioctl_cmds.h>
 #include <rdma/ib_ucaps.h>
 #include "macsec.h"
-#include "data_direct.h"
 #include "dmah.h"
 
 #define UVERBS_MODULE_NAME mlx5_ib
@@ -3962,20 +3961,9 @@ static int mlx5_ib_data_direct_init(struct mlx5_ib_dev *dev)
 	if (!mlx5_data_direct_supported(dev->mdev))
 		return 0;
 
-	ret = mlx5_data_direct_init(dev);
-	if (ret)
-		return ret;
-
 	INIT_LIST_HEAD(&dev->data_direct_mr_list);
 	dev->data_direct_nb.notifier_call = mlx5_ib_data_direct_event;
-	ret = mlx5_data_direct_register(dev, &dev->data_direct_nb);
-	if (ret)
-		goto err_register;
-
-	return ret;
-
-err_register:
-	mlx5_data_direct_cleanup(dev);
+	ret = mlx5_data_direct_register(dev->mdev, &dev->data_direct_nb);
 
 	return ret;
 }
@@ -3985,8 +3973,7 @@ static void mlx5_ib_data_direct_cleanup(struct mlx5_ib_dev *dev)
 	if (!mlx5_data_direct_supported(dev->mdev))
 		return;
 
-	mlx5_data_direct_unregister(dev, &dev->data_direct_nb);
-	mlx5_data_direct_cleanup(dev);
+	mlx5_data_direct_unregister(dev->mdev, &dev->data_direct_nb);
 }
 
 static int mlx5_ib_init_multiport_master(struct mlx5_ib_dev *dev)
@@ -5438,9 +5425,6 @@ static int __init mlx5_ib_init(void)
 	ret = mlx5r_rep_init();
 	if (ret)
 		goto rep_err;
-	ret = mlx5_data_direct_driver_register();
-	if (ret)
-		goto dd_err;
 	ret = auxiliary_driver_register(&mlx5r_mp_driver);
 	if (ret)
 		goto mp_err;
@@ -5453,8 +5437,6 @@ static int __init mlx5_ib_init(void)
 drv_err:
 	auxiliary_driver_unregister(&mlx5r_mp_driver);
 mp_err:
-	mlx5_data_direct_driver_unregister();
-dd_err:
 	mlx5r_rep_cleanup();
 rep_err:
 	mlx5_ib_qp_event_cleanup();
@@ -5466,7 +5448,6 @@ static int __init mlx5_ib_init(void)
 
 static void __exit mlx5_ib_cleanup(void)
 {
-	mlx5_data_direct_driver_unregister();
 	auxiliary_driver_unregister(&mlx5r_driver);
 	auxiliary_driver_unregister(&mlx5r_mp_driver);
 	mlx5r_rep_cleanup();
diff --git a/drivers/infiniband/hw/mlx5/mlx5_ib.h b/drivers/infiniband/hw/mlx5/mlx5_ib.h
index 6952170f472b..1a3b9cb79184 100644
--- a/drivers/infiniband/hw/mlx5/mlx5_ib.h
+++ b/drivers/infiniband/hw/mlx5/mlx5_ib.h
@@ -24,7 +24,7 @@
 #include <rdma/mlx5_user_ioctl_cmds.h>
 #include <rdma/mlx5_user_ioctl_verbs.h>
 
-#include "data_direct.h"
+#include <linux/mlx5/data_direct.h>
 #include "srq.h"
 #include "qp.h"
 #include "macsec.h"
@@ -1091,7 +1091,6 @@ struct mlx5_macsec {
 struct mlx5_ib_dev {
 	struct ib_device		ib_dev;
 	struct mlx5_core_dev		*mdev;
-	struct mlx5_data_direct		*data_direct;
 	/* Protects data_direct_mr_list and serializes mr
 	 * registration/deregistration with data direct device unbind.
 	 */
diff --git a/drivers/infiniband/hw/mlx5/mr.c b/drivers/infiniband/hw/mlx5/mr.c
index 5efe78cb34c1..7463b36a36b6 100644
--- a/drivers/infiniband/hw/mlx5/mr.c
+++ b/drivers/infiniband/hw/mlx5/mr.c
@@ -44,7 +44,6 @@
 #include "dm.h"
 #include "mlx5_ib.h"
 #include "umr.h"
-#include "data_direct.h"
 #include "dmah.h"
 
 static int mkey_max_umr_order(struct mlx5_ib_dev *dev)
@@ -587,7 +586,7 @@ static struct mlx5_ib_mr *reg_create(struct ib_pd *pd, struct ib_umem *umem,
 				      populate ? pd : dev->umrc.pd);
 	/* In case a data direct flow, overwrite the pdn field by its internal kernel PD */
 	if (umem->is_dmabuf && ksm_mode)
-		MLX5_SET(mkc, mkc, pd, dev->data_direct->pdn);
+		MLX5_SET(mkc, mkc, pd, dev->mdev->data_direct->pdn);
 
 	MLX5_SET(mkc, mkc, free, !populate);
 	MLX5_SET(mkc, mkc, access_mode_1_0, access_mode);
@@ -997,7 +996,7 @@ reg_user_mr_dmabuf_by_data_direct(struct ib_pd *pd, u64 offset,
 		return ERR_PTR(-EOPNOTSUPP);
 
 	mutex_lock(&dev->data_direct_lock);
-	data_direct_dev = mlx5_data_direct_get_dev(dev->data_direct);
+	data_direct_dev = mlx5_data_direct_get_dev(dev->mdev->data_direct);
 	if (!data_direct_dev) {
 		ret = -EINVAL;
 		goto end;
@@ -1006,7 +1005,7 @@ reg_user_mr_dmabuf_by_data_direct(struct ib_pd *pd, u64 offset,
 	/* If no device's 'data direct mkey' with RO flags exists
 	 * mask it out accordingly.
 	 */
-	if (!dev->data_direct->mkey_ro_valid)
+	if (!dev->mdev->data_direct->mkey_ro_valid)
 		access_flags &= ~IB_ACCESS_RELAXED_ORDERING;
 	crossed_mr = reg_user_mr_dmabuf(pd, &data_direct_dev->pdev->dev,
 					offset, length, virt_addr, fd,
diff --git a/drivers/infiniband/hw/mlx5/odp.c b/drivers/infiniband/hw/mlx5/odp.c
index 7b9f73b76d0d..2e9a704acc44 100644
--- a/drivers/infiniband/hw/mlx5/odp.c
+++ b/drivers/infiniband/hw/mlx5/odp.c
@@ -952,7 +952,7 @@ int mlx5_ib_init_dmabuf_mr(struct mlx5_ib_mr *mr, struct ib_pd *pd)
 	int ret;
 
 	if (mr->data_direct)
-		pdn = dev->data_direct->pdn;
+		pdn = dev->mdev->data_direct->pdn;
 	else
 		pdn = to_mpd(pd)->pdn;
 
diff --git a/drivers/infiniband/hw/mlx5/std_types.c b/drivers/infiniband/hw/mlx5/std_types.c
index 3828100fdc6c..b705b88dc917 100644
--- a/drivers/infiniband/hw/mlx5/std_types.c
+++ b/drivers/infiniband/hw/mlx5/std_types.c
@@ -10,7 +10,6 @@
 #include <linux/mlx5/eswitch.h>
 #include <linux/mlx5/vport.h>
 #include "mlx5_ib.h"
-#include "data_direct.h"
 
 #define UVERBS_MODULE_NAME mlx5_ib
 #include <rdma/uverbs_named_ioctl.h>
@@ -203,7 +202,7 @@ static int UVERBS_HANDLER(MLX5_IB_METHOD_GET_DATA_DIRECT_SYSFS_PATH)(
 		return PTR_ERR(c);
 	dev = to_mdev(c->ibucontext.device);
 	mutex_lock(&dev->data_direct_lock);
-	data_direct_dev = mlx5_data_direct_get_dev(dev->data_direct);
+	data_direct_dev = mlx5_data_direct_get_dev(dev->mdev->data_direct);
 	if (!data_direct_dev) {
 		ret = -ENODEV;
 		goto end;
diff --git a/drivers/infiniband/hw/mlx5/umr.c b/drivers/infiniband/hw/mlx5/umr.c
index 1ca83d2d8a0b..a3bb8d03fd35 100644
--- a/drivers/infiniband/hw/mlx5/umr.c
+++ b/drivers/infiniband/hw/mlx5/umr.c
@@ -760,7 +760,7 @@ static int _mlx5r_umr_update_mr_pas(struct mlx5_ib_mr *mr, unsigned int flags,
 		}
 
 		if (dd) {
-			struct mlx5_data_direct *mdd = dev->data_direct;
+			struct mlx5_data_direct *mdd = dev->mdev->data_direct;
 
 			cur_ksm->va = cpu_to_be64(rdma_block_iter_dma_address(&biter));
 			if (mr->access_flags & IB_ACCESS_RELAXED_ORDERING &&
@@ -813,7 +813,7 @@ int mlx5r_umr_update_data_direct_ksm_pas_range(struct mlx5_ib_mr *mr,
 		return -EINVAL;
 
 	return _mlx5r_umr_update_mr_pas(mr, flags,
-					mr_to_mdev(mr)->data_direct->pdn,
+					mr_to_mdev(mr)->mdev->data_direct->pdn,
 					true, start_block, nblocks);
 }
 
diff --git a/drivers/net/ethernet/mellanox/mlx5/core/Makefile b/drivers/net/ethernet/mellanox/mlx5/core/Makefile
index 19e50f0d55af..a1871259210d 100644
--- a/drivers/net/ethernet/mellanox/mlx5/core/Makefile
+++ b/drivers/net/ethernet/mellanox/mlx5/core/Makefile
@@ -18,7 +18,7 @@ mlx5_core-y :=	main.o cmd.o debugfs.o fw.o eq.o uar.o pagealloc.o \
 		lib/devcom.o lib/pci_vsc.o lib/dm.o lib/fs_ttc.o diag/fs_tracepoint.o \
 		diag/fw_tracer.o diag/crdump.o devlink.o sh_devlink.o diag/rsc_dump.o \
 		diag/reporter_vnic.o fw_reset.o qos.o lib/tout.o lib/aso.o wc.o fs_pool.o \
-		lib/nv_param.o
+		lib/nv_param.o data_direct.o
 
 #
 # Netdev basic
diff --git a/drivers/infiniband/hw/mlx5/data_direct.c b/drivers/net/ethernet/mellanox/mlx5/core/data_direct.c
similarity index 79%
rename from drivers/infiniband/hw/mlx5/data_direct.c
rename to drivers/net/ethernet/mellanox/mlx5/core/data_direct.c
index 3878c4aebb2c..67717110feab 100644
--- a/drivers/infiniband/hw/mlx5/data_direct.c
+++ b/drivers/net/ethernet/mellanox/mlx5/core/data_direct.c
@@ -3,11 +3,14 @@
  * Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved
  */
 
-#include "mlx5_ib.h"
+#include <linux/list.h>
+#include <linux/mlx5/driver.h>
+#include <linux/mlx5/data_direct.h>
+#include <linux/mlx5/mlx5_ifc.h>
+#include <linux/pci.h>
+#include <linux/slab.h>
 
-#include <linux/notifier.h>
-
-#include "data_direct.h"
+#include "mlx5_core.h"
 
 static LIST_HEAD(mlx5_data_direct_dev_list);
 static LIST_HEAD(mlx5_data_direct_reg_list);
@@ -103,10 +106,9 @@ static int mlx5_data_direct_set_dma_caps(struct pci_dev *pdev)
 	return 0;
 }
 
-static int mlx5_data_direct_create_resources(struct mlx5_ib_dev *dev)
+static int mlx5_data_direct_create_resources(struct mlx5_core_dev *mdev)
 {
 	int inlen = MLX5_ST_SZ_BYTES(create_mkey_in);
-	struct mlx5_core_dev *mdev = dev->mdev;
 	bool ro_supp = false;
 	void *mkc;
 	u32 mkey;
@@ -139,16 +141,16 @@ static int mlx5_data_direct_create_resources(struct mlx5_ib_dev *dev)
 	if (err)
 		goto err_mkey;
 
-	dev->data_direct->mkey = mkey;
-	dev->data_direct->pdn = pdn;
+	mdev->data_direct->mkey = mkey;
+	mdev->data_direct->pdn = pdn;
 
 	/* create another mkey with RO support */
-	if (MLX5_CAP_GEN(dev->mdev, relaxed_ordering_write)) {
+	if (MLX5_CAP_GEN(mdev, relaxed_ordering_write)) {
 		MLX5_SET(mkc, mkc, relaxed_ordering_write, 1);
 		ro_supp = true;
 	}
 
-	if (MLX5_CAP_GEN(dev->mdev, relaxed_ordering_read)) {
+	if (MLX5_CAP_GEN(mdev, relaxed_ordering_read)) {
 		MLX5_SET(mkc, mkc, relaxed_ordering_read, 1);
 		ro_supp = true;
 	}
@@ -157,8 +159,8 @@ static int mlx5_data_direct_create_resources(struct mlx5_ib_dev *dev)
 		err = mlx5_core_create_mkey(mdev, &mkey, in, inlen);
 		/* RO is defined as best effort */
 		if (!err) {
-			dev->data_direct->mkey_ro = mkey;
-			dev->data_direct->mkey_ro_valid = true;
+			mdev->data_direct->mkey_ro = mkey;
+			mdev->data_direct->mkey_ro_valid = true;
 		}
 	}
 
@@ -172,13 +174,15 @@ static int mlx5_data_direct_create_resources(struct mlx5_ib_dev *dev)
 	return err;
 }
 
-static void mlx5_data_direct_free_resources(struct mlx5_ib_dev *dev)
+static void mlx5_data_direct_free_resources(struct mlx5_core_dev *mdev)
 {
-	if (dev->data_direct->mkey_ro_valid)
-		mlx5_core_destroy_mkey(dev->mdev, dev->data_direct->mkey_ro);
+	struct mlx5_data_direct *data_direct = mdev->data_direct;
+
+	if (data_direct->mkey_ro_valid)
+		mlx5_core_destroy_mkey(mdev, data_direct->mkey_ro);
 
-	mlx5_core_destroy_mkey(dev->mdev, dev->data_direct->mkey);
-	mlx5_core_dealloc_pd(dev->mdev, dev->data_direct->pdn);
+	mlx5_core_destroy_mkey(mdev, data_direct->mkey);
+	mlx5_core_dealloc_pd(mdev, data_direct->pdn);
 }
 
 static void mlx5_data_direct_bind(struct mlx5_data_direct_registration *reg,
@@ -195,13 +199,13 @@ mlx5_data_direct_do_unbind(struct mlx5_data_direct_registration *reg)
 				     NULL);
 }
 
-int mlx5_data_direct_init(struct mlx5_ib_dev *ibdev)
+int mlx5_data_direct_init(struct mlx5_core_dev *mdev)
 {
 	struct mlx5_data_direct_registration *reg;
 	struct mlx5_data_direct_dev *dev;
 	int err;
 
-	if (!mlx5_data_direct_supported(ibdev->mdev))
+	if (!mlx5_data_direct_supported(mdev))
 		return 0;
 
 	reg = kzalloc_obj(*reg);
@@ -210,17 +214,18 @@ int mlx5_data_direct_init(struct mlx5_ib_dev *ibdev)
 
 	BLOCKING_INIT_NOTIFIER_HEAD(&reg->users);
 
-	err = mlx5_data_direct_query_vuid(ibdev->mdev, reg->vuid);
+	err = mlx5_data_direct_query_vuid(mdev, reg->vuid);
 	if (err) {
-		mlx5_ib_warn(ibdev, "Failed to query VUID, disabling data direct, err=%d\n",
-			     err);
+		mlx5_core_warn(mdev,
+			       "Failed to query VUID, disabling data direct, err=%d\n",
+			       err);
 		kfree(reg);
-		return err;
+		return -EINVAL;
 	}
 
-	ibdev->data_direct = &reg->dd;
+	mdev->data_direct = &reg->dd;
 
-	err = mlx5_data_direct_create_resources(ibdev);
+	err = mlx5_data_direct_create_resources(mdev);
 	if (err)
 		goto err_resources;
 
@@ -233,63 +238,63 @@ int mlx5_data_direct_init(struct mlx5_ib_dev *ibdev)
 	}
 
 	/* Add the registration to its global list, to be used upon bind/unbind
-	 * of its affiliated data direct device
+	 * of its affiliated data direct device.
 	 */
 	list_add_tail(&reg->list, &mlx5_data_direct_reg_list);
 	mutex_unlock(&mlx5_data_direct_mutex);
 	return 0;
 
 err_resources:
-	ibdev->data_direct = NULL;
+	mdev->data_direct = NULL;
 	kfree(reg);
 	return err;
 }
 
-void mlx5_data_direct_cleanup(struct mlx5_ib_dev *ibdev)
+void mlx5_data_direct_cleanup(struct mlx5_core_dev *mdev)
 {
 	struct mlx5_data_direct_registration *reg;
 
-	if (!mlx5_data_direct_supported(ibdev->mdev))
+	if (!mlx5_data_direct_supported(mdev))
 		return;
 
-	reg = container_of(ibdev->data_direct,
-			    struct mlx5_data_direct_registration, dd);
+	reg = container_of(mdev->data_direct,
+			   struct mlx5_data_direct_registration, dd);
 	mutex_lock(&mlx5_data_direct_mutex);
 	list_del(&reg->list);
 	mlx5_data_direct_do_unbind(reg);
 	mutex_unlock(&mlx5_data_direct_mutex);
 
-	mlx5_data_direct_free_resources(ibdev);
-	ibdev->data_direct = NULL;
+	mlx5_data_direct_free_resources(mdev);
+	mdev->data_direct = NULL;
 	kfree(reg);
 }
 
-int mlx5_data_direct_register(struct mlx5_ib_dev *ibdev,
+int mlx5_data_direct_register(struct mlx5_core_dev *mdev,
 			      struct notifier_block *nb)
 {
 	struct mlx5_data_direct_registration *reg;
 
-	if (!mlx5_data_direct_supported(ibdev->mdev))
+	if (!mlx5_data_direct_supported(mdev))
 		return 0;
 
-	reg = container_of(ibdev->data_direct,
-			    struct mlx5_data_direct_registration, dd);
+	reg = container_of(mdev->data_direct,
+			   struct mlx5_data_direct_registration, dd);
 	blocking_notifier_chain_register(&reg->users, nb);
 
 	return 0;
 }
 EXPORT_SYMBOL_GPL(mlx5_data_direct_register);
 
-void mlx5_data_direct_unregister(struct mlx5_ib_dev *ibdev,
+void mlx5_data_direct_unregister(struct mlx5_core_dev *mdev,
 				 struct notifier_block *nb)
 {
 	struct mlx5_data_direct_registration *reg;
 
-	if (!mlx5_data_direct_supported(ibdev->mdev))
+	if (!mlx5_data_direct_supported(mdev))
 		return;
 
-	reg = container_of(ibdev->data_direct,
-			    struct mlx5_data_direct_registration, dd);
+	reg = container_of(mdev->data_direct,
+			   struct mlx5_data_direct_registration, dd);
 	blocking_notifier_chain_unregister(&reg->users, nb);
 }
 EXPORT_SYMBOL_GPL(mlx5_data_direct_unregister);
@@ -304,8 +309,8 @@ static void mlx5_data_direct_dev_reg(struct mlx5_data_direct_dev *dev)
 			mlx5_data_direct_bind(reg, dev);
 	}
 
-	/* Add the data direct device to the global list, further IB devices may
-	 * use it later as well
+	/* Add the data direct device to the global list, further mlx5 devices
+	 * may use it later as well.
 	 */
 	list_add_tail(&dev->list, &mlx5_data_direct_dev_list);
 	mutex_unlock(&mlx5_data_direct_mutex);
diff --git a/drivers/net/ethernet/mellanox/mlx5/core/main.c b/drivers/net/ethernet/mellanox/mlx5/core/main.c
index 643b4aac2033..34f0375b6851 100644
--- a/drivers/net/ethernet/mellanox/mlx5/core/main.c
+++ b/drivers/net/ethernet/mellanox/mlx5/core/main.c
@@ -40,6 +40,7 @@
 #include <linux/interrupt.h>
 #include <linux/delay.h>
 #include <linux/mlx5/driver.h>
+#include <linux/mlx5/data_direct.h>
 #include <linux/mlx5/cq.h>
 #include <linux/mlx5/qp.h>
 #include <linux/debugfs.h>
@@ -1339,8 +1340,14 @@ static int mlx5_load(struct mlx5_core_dev *dev)
 	if (err)
 		goto err_traps_reg;
 
+	err = mlx5_data_direct_init(dev);
+	if (err)
+		goto err_data_direct;
+
 	return 0;
 
+err_data_direct:
+	mlx5_devlink_traps_unregister(priv_to_devlink(dev));
 err_traps_reg:
 	mlx5_sf_dev_table_destroy(dev);
 	mlx5_sriov_detach(dev);
@@ -1371,6 +1378,7 @@ static int mlx5_load(struct mlx5_core_dev *dev)
 
 static void mlx5_unload(struct mlx5_core_dev *dev)
 {
+	mlx5_data_direct_cleanup(dev);
 	mlx5_devlink_traps_unregister(priv_to_devlink(dev));
 	mlx5_vhca_event_stop(dev);
 	mlx5_sf_dev_table_destroy(dev);
@@ -2360,6 +2368,10 @@ static int __init mlx5_init(void)
 	if (err)
 		goto err_sf;
 
+	err = mlx5_data_direct_driver_register();
+	if (err)
+		goto err_dd;
+
 	err = pci_register_driver(&mlx5_core_driver);
 	if (err)
 		goto err_pci;
@@ -2367,6 +2379,8 @@ static int __init mlx5_init(void)
 	return 0;
 
 err_pci:
+	mlx5_data_direct_driver_unregister();
+err_dd:
 	mlx5_sf_driver_unregister();
 err_sf:
 	mlx5e_cleanup();
@@ -2378,6 +2392,7 @@ static int __init mlx5_init(void)
 static void __exit mlx5_cleanup(void)
 {
 	pci_unregister_driver(&mlx5_core_driver);
+	mlx5_data_direct_driver_unregister();
 	mlx5_sf_driver_unregister();
 	mlx5e_cleanup();
 	mlx5_unregister_debugfs();
diff --git a/drivers/net/ethernet/mellanox/mlx5/core/mlx5_core.h b/drivers/net/ethernet/mellanox/mlx5/core/mlx5_core.h
index 09e669f83dba..fd93530a3938 100644
--- a/drivers/net/ethernet/mellanox/mlx5/core/mlx5_core.h
+++ b/drivers/net/ethernet/mellanox/mlx5/core/mlx5_core.h
@@ -530,4 +530,9 @@ static inline bool mlx5_pcie_cong_event_supported(struct mlx5_core_dev *dev)
 
 	return true;
 }
+
+int mlx5_data_direct_init(struct mlx5_core_dev *mdev);
+void mlx5_data_direct_cleanup(struct mlx5_core_dev *mdev);
+int mlx5_data_direct_driver_register(void);
+void mlx5_data_direct_driver_unregister(void);
 #endif /* __MLX5_CORE_H__ */
diff --git a/include/linux/mlx5/data_direct.h b/include/linux/mlx5/data_direct.h
index eaa2dc6bc607..4cff5ff5489d 100644
--- a/include/linux/mlx5/data_direct.h
+++ b/include/linux/mlx5/data_direct.h
@@ -7,11 +7,23 @@
 #define _MLX5_DATA_DIRECT_H
 
 #include <linux/compiler.h>
+#include <linux/list.h>
 
+struct device;
 struct mlx5_core_dev;
-struct mlx5_data_direct_dev;
-struct mlx5_ib_dev;
 struct notifier_block;
+struct pci_dev;
+
+enum mlx5_data_direct_event {
+	MLX5_DATA_DIRECT_UNBIND,
+};
+
+struct mlx5_data_direct_dev {
+	struct device *device;
+	struct pci_dev *pdev;
+	char *vuid;
+	struct list_head list;
+};
 
 struct mlx5_data_direct {
 	struct mlx5_data_direct_dev *dev;
@@ -27,12 +39,9 @@ mlx5_data_direct_get_dev(struct mlx5_data_direct *dd)
 	return dd ? READ_ONCE(dd->dev) : NULL;
 }
 
-int mlx5_data_direct_init(struct mlx5_ib_dev *ibdev);
-void mlx5_data_direct_cleanup(struct mlx5_ib_dev *ibdev);
-
-int mlx5_data_direct_register(struct mlx5_ib_dev *ibdev,
+int mlx5_data_direct_register(struct mlx5_core_dev *mdev,
 			      struct notifier_block *nb);
-void mlx5_data_direct_unregister(struct mlx5_ib_dev *ibdev,
+void mlx5_data_direct_unregister(struct mlx5_core_dev *mdev,
 				 struct notifier_block *nb);
 
 #endif
diff --git a/include/linux/mlx5/driver.h b/include/linux/mlx5/driver.h
index ee29f95ead18..ed348a82576e 100644
--- a/include/linux/mlx5/driver.h
+++ b/include/linux/mlx5/driver.h
@@ -707,6 +707,7 @@ struct mlx5_rsvd_gids {
 
 struct mlx5_clock;
 struct mlx5_clock_dev_state;
+struct mlx5_data_direct;
 struct mlx5_dm;
 struct mlx5_fw_tracer;
 struct mlx5_vxlan;
@@ -801,6 +802,7 @@ struct mlx5_core_dev {
 	/* sync write combining state */
 	struct mutex wc_state_lock;
 	struct devlink *shd;
+	struct mlx5_data_direct *data_direct;
 };
 
 struct mlx5_db {
-- 
2.44.0