From b297f1b2d32518a67fe2235b9697340c9e99e393 Mon Sep 17 00:00:00 2001
From: David Nadlinger <code@klickverbot.at>
Date: Sun, 8 Oct 2023 14:10:00 +0100
Subject: [PATCH] compiler: Catch escaping numpy.{array, full, transpose}()
 results

Function calls in general can still be used to hide escaping
allocations from the compiler (issue #1497), but these calls in
particular always allocate, so we can easily and accurately handle
them.
---
 artiq/compiler/validators/escape.py            | 16 ++++++++++++++--
 artiq/test/lit/escape/error_numpy_array.py     | 15 +++++++++++++++
 artiq/test/lit/escape/error_numpy_full.py      | 16 ++++++++++++++++
 artiq/test/lit/escape/error_numpy_transpose.py | 17 +++++++++++++++++
 doc/manual/compiler.rst                        | 18 ------------------
 5 files changed, 62 insertions(+), 20 deletions(-)
 create mode 100644 artiq/test/lit/escape/error_numpy_array.py
 create mode 100644 artiq/test/lit/escape/error_numpy_full.py
 create mode 100644 artiq/test/lit/escape/error_numpy_transpose.py

diff --git a/artiq/compiler/validators/escape.py b/artiq/compiler/validators/escape.py
index c6ae597041..7d23958301 100644
--- a/artiq/compiler/validators/escape.py
+++ b/artiq/compiler/validators/escape.py
@@ -102,8 +102,20 @@ def visit_CallT(self, node):
         if types.is_external_function(node.func.type, "cache_get"):
             # The cache is borrow checked dynamically
             return Global()
-        else:
-            self.visit_sometimes_allocating(node)
+
+        if (types.is_builtin_function(node.func.type, "array")
+              or types.is_builtin_function(node.func.type, "make_array")
+              or types.is_builtin_function(node.func.type, "numpy.transpose")):
+            # While lifetime tracking across function calls in general is currently
+            # broken (see below), these special builtins that allocate an array on
+            # the stack of the caller _always_ allocate regardless of the parameters,
+            # and we can thus handle them without running into the precision issue
+            # mentioned in commit ae999db.
+            return self.visit_allocating(node)
+
+        # FIXME: Return statement missing here, but see m-labs/artiq#1497 and
+        # commit ae999db.
+        self.visit_sometimes_allocating(node)
 
     # Value lives as long as the object/container, if it's mutable,
     # or else forever
diff --git a/artiq/test/lit/escape/error_numpy_array.py b/artiq/test/lit/escape/error_numpy_array.py
new file mode 100644
index 0000000000..1ebdeda818
--- /dev/null
+++ b/artiq/test/lit/escape/error_numpy_array.py
@@ -0,0 +1,15 @@
+# RUN: %python -m artiq.compiler.testbench.embedding +diag %s 2>%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+from artiq.experiment import *
+import numpy as np
+
+@kernel
+def a():
+    # CHECK-L: ${LINE:+2}: error: cannot return an allocated value that does not live forever
+    # CHECK-L: ${LINE:+1}: note: ... to this point
+    return np.array([0, 1])
+
+@kernel
+def entrypoint():
+    a()
diff --git a/artiq/test/lit/escape/error_numpy_full.py b/artiq/test/lit/escape/error_numpy_full.py
new file mode 100644
index 0000000000..66158d8cad
--- /dev/null
+++ b/artiq/test/lit/escape/error_numpy_full.py
@@ -0,0 +1,16 @@
+# RUN: %python -m artiq.compiler.testbench.embedding +diag %s 2>%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+from artiq.experiment import *
+import numpy as np
+
+@kernel
+def a():
+    # CHECK-L: ${LINE:+2}: error: cannot return an allocated value that does not live forever
+    # CHECK-L: ${LINE:+1}: note: ... to this point
+    return np.full(10, 42.0)
+
+
+@kernel
+def entrypoint():
+    a()
diff --git a/artiq/test/lit/escape/error_numpy_transpose.py b/artiq/test/lit/escape/error_numpy_transpose.py
new file mode 100644
index 0000000000..e1dc32d51d
--- /dev/null
+++ b/artiq/test/lit/escape/error_numpy_transpose.py
@@ -0,0 +1,17 @@
+# RUN: %python -m artiq.compiler.testbench.embedding +diag %s 2>%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+from artiq.experiment import *
+import numpy as np
+
+data = np.array([[0, 1], [2, 3]])
+
+@kernel
+def a():
+    # CHECK-L: ${LINE:+2}: error: cannot return an allocated value that does not live forever
+    # CHECK-L: ${LINE:+1}: note: ... to this point
+    return np.transpose(data)
+
+@kernel
+def entrypoint():
+    a()
diff --git a/doc/manual/compiler.rst b/doc/manual/compiler.rst
index e4923dfa93..88b57d8147 100644
--- a/doc/manual/compiler.rst
+++ b/doc/manual/compiler.rst
@@ -95,24 +95,6 @@ tracked across function calls (see `#1497 <https://github.com/m-labs/artiq/issue
         def run(self):
             # results in memory corruption
             return func([1, 2, 3])
-or if the return value is obfuscated by an if-statement like here: ::
-
-    class ProblemReturn2(EnvExperiment):
-        def build(self):
-            self.setattr_device("core")
-
-        @kernel
-        def meth(self):
-            # if statement for obfuscation
-            if self.core.get_rtio_counter_mu() % 2:
-                return np.array([1,2,3])
-            else:
-                return np.array([4,5,6])
-
-        @kernel
-        def run(self):
-            # also results in memory corrption
-            return self.meth()
 
 This results in memory corruption at runtime.