From 01428ab9f61ea916d1a9d7551d6d7d297c936d50 Mon Sep 17 00:00:00 2001
From: "Wang, Mingjie1" <mingjie1.wang@intel.com>
Date: Tue, 28 Mar 2023 12:47:04 -0500
Subject: [PATCH] Updated mk_alloc to support Numba-Dpex compute follows data.

---
 numba/core/ir_utils.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/numba/core/ir_utils.py b/numba/core/ir_utils.py
index 8586c8f79d1..f5041e66393 100644
--- a/numba/core/ir_utils.py
+++ b/numba/core/ir_utils.py
@@ -91,6 +91,20 @@ def mk_alloc(typingctx, typemap, calltypes, lhs, size_var, dtype, scope, loc,
             out.append(tuple_assign)
             size_var = tuple_var
             size_typ = types.containers.UniTuple(types.intp, ndims)
+    if hasattr(lhs_typ, "__allocate__"):
+        return lhs_typ.__allocate__(
+            typingctx,
+            typemap,
+            calltypes,
+            lhs,
+            size_var,
+            dtype,
+            scope,
+            loc,
+            lhs_typ,
+            size_typ,
+            out,
+        )
     # g_np_var = Global(numpy)
     g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
     if typemap: