From 8d7194941ecc0fd0e52f38c3f44bbf476a903d23 Mon Sep 17 00:00:00 2001
From: mwojcik <mw@m-labs.hk>
Date: Thu, 5 Oct 2023 14:36:09 +0800
Subject: [PATCH] tests: add lit tests for subkernels

---
 .../lit/embedding/error_subkernel_annot.py    | 15 +++++++++++
 .../embedding/error_subkernel_annot_return.py | 15 +++++++++++
 artiq/test/lit/embedding/subkernel_no_arg.py  | 18 +++++++++++++
 artiq/test/lit/embedding/subkernel_return.py  | 22 ++++++++++++++++
 .../lit/embedding/subkernel_return_none.py    | 22 ++++++++++++++++
 artiq/test/lit/embedding/subkernel_self.py    | 25 +++++++++++++++++++
 .../test/lit/embedding/subkernel_self_args.py | 25 +++++++++++++++++++
 .../test/lit/embedding/subkernel_with_arg.py  | 18 +++++++++++++
 .../lit/embedding/subkernel_with_opt_arg.py   | 21 ++++++++++++++++
 9 files changed, 181 insertions(+)
 create mode 100644 artiq/test/lit/embedding/error_subkernel_annot.py
 create mode 100644 artiq/test/lit/embedding/error_subkernel_annot_return.py
 create mode 100644 artiq/test/lit/embedding/subkernel_no_arg.py
 create mode 100644 artiq/test/lit/embedding/subkernel_return.py
 create mode 100644 artiq/test/lit/embedding/subkernel_return_none.py
 create mode 100644 artiq/test/lit/embedding/subkernel_self.py
 create mode 100644 artiq/test/lit/embedding/subkernel_self_args.py
 create mode 100644 artiq/test/lit/embedding/subkernel_with_arg.py
 create mode 100644 artiq/test/lit/embedding/subkernel_with_opt_arg.py

diff --git a/artiq/test/lit/embedding/error_subkernel_annot.py b/artiq/test/lit/embedding/error_subkernel_annot.py
new file mode 100644
index 000000000..3f4bc1c5c
--- /dev/null
+++ b/artiq/test/lit/embedding/error_subkernel_annot.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.language.core import *
+from artiq.language.types import *
+
+# CHECK-L: ${LINE:+2}: error: type annotation for argument 'x', '1', is not an ARTIQ type
+@subkernel(destination=1)
+def foo(x: 1) -> TNone:
+    pass
+
+@kernel
+def entrypoint():
+    # CHECK-L: ${LINE:+1}: note: in subkernel call here
+    foo()
diff --git a/artiq/test/lit/embedding/error_subkernel_annot_return.py b/artiq/test/lit/embedding/error_subkernel_annot_return.py
new file mode 100644
index 000000000..51977ae00
--- /dev/null
+++ b/artiq/test/lit/embedding/error_subkernel_annot_return.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.language.core import *
+from artiq.language.types import *
+
+# CHECK-L: ${LINE:+2}: error: type annotation for return type, '1', is not an ARTIQ type
+@subkernel(destination=1)
+def foo() -> 1:
+    pass
+
+@kernel
+def entrypoint():
+    # CHECK-L: ${LINE:+1}: note: in subkernel call here
+    foo()
diff --git a/artiq/test/lit/embedding/subkernel_no_arg.py b/artiq/test/lit/embedding/subkernel_no_arg.py
new file mode 100644
index 000000000..11d9c793d
--- /dev/null
+++ b/artiq/test/lit/embedding/subkernel_no_arg.py
@@ -0,0 +1,18 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+@kernel
+def entrypoint():
+    # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+    # CHECK-NOT: call void @subkernel_send_message\(.*\), !dbg !.
+    no_arg()
+
+
+# CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
+# CHECK-NOT-L: declare void @subkernel_send_message(i32, { i8*, i32 }*, i8**) local_unnamed_addr
+@subkernel(destination=1)
+def no_arg() -> TStr:
+    pass
diff --git a/artiq/test/lit/embedding/subkernel_return.py b/artiq/test/lit/embedding/subkernel_return.py
new file mode 100644
index 000000000..4845e24ba
--- /dev/null
+++ b/artiq/test/lit/embedding/subkernel_return.py
@@ -0,0 +1,22 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+@kernel
+def entrypoint():
+    # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+    # CHECK-NOT: call void @subkernel_send_message\(.*\), !dbg !.
+    returning()
+    # CHECK: call i8 @subkernel_await_message\(i32 1, i64 10000, i8 1, i8 1\), !dbg !.
+    # CHECK: call void @subkernel_await_finish\(i32 1, i64 10000\), !dbg !.
+    subkernel_await(returning)
+
+# CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
+# CHECK-NOT-L: declare void @subkernel_send_message(i32, i8, { i8*, i32 }*, i8**) local_unnamed_addr
+# CHECK-L: declare i8 @subkernel_await_message(i32, i64, i8, i8) local_unnamed_addr
+# CHECK-L: declare void @subkernel_await_finish(i32, i64) local_unnamed_addr
+@subkernel(destination=1)
+def returning() -> TInt32:
+    return 1
diff --git a/artiq/test/lit/embedding/subkernel_return_none.py b/artiq/test/lit/embedding/subkernel_return_none.py
new file mode 100644
index 000000000..353b15c3e
--- /dev/null
+++ b/artiq/test/lit/embedding/subkernel_return_none.py
@@ -0,0 +1,22 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+@kernel
+def entrypoint():
+    # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+    # CHECK-NOT: call void @subkernel_send_message\(.*\), !dbg !.
+    returning_none()
+    # CHECK: call void @subkernel_await_finish\(i32 1, i64 10000\), !dbg !.
+    # CHECK-NOT: call void @subkernel_await_message\(i32 1, i64 10000\), !dbg !.
+    subkernel_await(returning_none)
+
+# CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
+# CHECK-NOT-L: declare void @subkernel_send_message(i32, { i8*, i32 }*, i8**) local_unnamed_addr
+# CHECK-L: declare void @subkernel_await_finish(i32, i64) local_unnamed_addr
+# CHECK-NOT-L: declare void @subkernel_await_message(i32, i64) local_unnamed_addr
+@subkernel(destination=1)
+def returning_none() -> TNone:
+    pass
diff --git a/artiq/test/lit/embedding/subkernel_self.py b/artiq/test/lit/embedding/subkernel_self.py
new file mode 100644
index 000000000..8e702bc02
--- /dev/null
+++ b/artiq/test/lit/embedding/subkernel_self.py
@@ -0,0 +1,25 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+class A:
+    @subkernel(destination=1)
+    def sk(self):
+        pass
+
+    @kernel
+    def kernel_entrypoint(self):
+        # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+        # CHECK-NOT: call void @subkernel_send_message\(.*\), !dbg !.
+        self.sk()
+
+a = A()
+
+@kernel
+def entrypoint():
+    a.kernel_entrypoint()
+
+# CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
+# CHECK-NOT-L: declare void @subkernel_send_message(i32, i8, { i8*, i32 }*, i8**) local_unnamed_addr
diff --git a/artiq/test/lit/embedding/subkernel_self_args.py b/artiq/test/lit/embedding/subkernel_self_args.py
new file mode 100644
index 000000000..57969398c
--- /dev/null
+++ b/artiq/test/lit/embedding/subkernel_self_args.py
@@ -0,0 +1,25 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+class A:
+    @subkernel(destination=1)
+    def sk(self, a):
+        pass
+
+    @kernel
+    def kernel_entrypoint(self):
+        # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+        # CHECK: call void @subkernel_send_message\(i32 1, i8 1, .*\), !dbg !.
+        self.sk(1)
+
+a = A()
+
+@kernel
+def entrypoint():
+    a.kernel_entrypoint()
+
+# CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
+# CHECK-L: declare void @subkernel_send_message(i32, i8, { i8*, i32 }*, i8**) local_unnamed_addr
diff --git a/artiq/test/lit/embedding/subkernel_with_arg.py b/artiq/test/lit/embedding/subkernel_with_arg.py
new file mode 100644
index 000000000..17d80dce7
--- /dev/null
+++ b/artiq/test/lit/embedding/subkernel_with_arg.py
@@ -0,0 +1,18 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+@kernel
+def entrypoint():
+    # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+    # CHECK: call void @subkernel_send_message\(i32 ., i8 1, .*\), !dbg !.
+    accept_arg(1)
+
+
+# CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
+# CHECK-L: declare void @subkernel_send_message(i32, i8, { i8*, i32 }*, i8**) local_unnamed_addr
+@subkernel(destination=1)
+def accept_arg(arg: TInt32) -> TNone:
+    pass
diff --git a/artiq/test/lit/embedding/subkernel_with_opt_arg.py b/artiq/test/lit/embedding/subkernel_with_opt_arg.py
new file mode 100644
index 000000000..1821fec01
--- /dev/null
+++ b/artiq/test/lit/embedding/subkernel_with_opt_arg.py
@@ -0,0 +1,21 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+@kernel
+def entrypoint():
+    # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+    # CHECK: call void @subkernel_send_message\(i32 ., i8 1, .*\), !dbg !.
+    accept_arg(1)
+    # CHECK: call void @subkernel_load_run\(i32 1, i1 true\), !dbg !.
+    # CHECK: call void @subkernel_send_message\(i32 ., i8 2, .*\), !dbg !.
+    accept_arg(1, 2)
+
+
+# CHECK-L: declare void @subkernel_load_run(i32, i1) local_unnamed_addr
+# CHECK-L: declare void @subkernel_send_message(i32, i8, { i8*, i32 }*, i8**) local_unnamed_addr
+@subkernel(destination=1)
+def accept_arg(arg_a, arg_b=5) -> TNone:
+    pass