Updated mk_alloc to support Numba-Dpex compute follows data.

Deezy33-creator · Mar 28, 2023 · 01428ab · 01428ab
1 parent 24b7628
commit 01428ab
Showing 1 changed file with 14 additions and 0 deletions.
diff --git 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: