When hl_dio_get_iopath() succeeds, it calls hl_ctx_get() to acquire a
reference on the context. If the subsequent vzalloc() fails, the
function returns -ENOMEM without calling hl_dio_put_iopath(), leaking
the reference.

Fix this by jumping to the cleanup label on error, which will call
hl_dio_put_iopath() and safely handle the NULL io->bv.

Cc: stable@vger.kernel.org
Fixes: 8cbacc9a2703 ("accel/habanalabs: add NVMe Direct I/O (HLDIO) infrastructure")
Signed-off-by: Wentao Liang <vulab@iscas.ac.cn>
---
 drivers/accel/habanalabs/common/hldio.c | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/drivers/accel/habanalabs/common/hldio.c b/drivers/accel/habanalabs/common/hldio.c
index c33c817a962a..d7dfa259bede 100644
--- a/drivers/accel/habanalabs/common/hldio.c
+++ b/drivers/accel/habanalabs/common/hldio.c
@@ -248,8 +248,10 @@ static ssize_t hl_direct_io(struct hl_device *hdev, struct hl_direct_io *io)
 	 * closest one.
 	 */
 	io->bv = vzalloc(npages * sizeof(struct bio_vec));
-	if (!io->bv)
+	if (!io->bv) {
+		hl_dio_put_iopath(io->f.ctx);
 		return -ENOMEM;
+	}
 
 	for (i = 0, device_va = io->device_va; i < npages ; ++i, device_va += PAGE_SIZE) {
 		io->bv[i].bv_page = hl_dio_va2page(hdev, io->f.ctx, device_va);
-- 
2.34.1