m-labs · Jul 7, 2016 · Jul 7, 2016
Showing with 74 additions and 2 deletions.

+9 −0 artiq/compiler/builtins.py

+21 −1 artiq/compiler/embedding.py

+16 −1 artiq/compiler/transforms/artiq_ir_generator.py

+17 −0 artiq/compiler/transforms/inferencer.py

+11 −0 artiq/test/coredevice/test_embedding.py
diff --git a/artiq/compiler/builtins.py b/artiq/compiler/builtins.py
@@ -125,6 +125,12 @@ def fn_bool():
 def fn_int():
     return types.TConstructor(TInt())
 
+def fn_int32():
+    return types.TConstructor(TInt32())
+
+def fn_int64():
+    return types.TConstructor(TInt64())
+
 def fn_float():
     return types.TConstructor(TFloat())
 
@@ -164,6 +170,9 @@ def fn_min():
 def fn_max():
     return types.TBuiltinFunction("max")
 
+def fn_make_array():
+    return types.TBuiltinFunction("make_array")
+
 def fn_print():
     return types.TBuiltinFunction("print")
 

diff --git a/artiq/compiler/embedding.py b/artiq/compiler/embedding.py
@@ -158,6 +158,22 @@ def quote(self, value):
             typ = builtins.TBool()
             return asttyped.NameConstantT(value=value, type=typ,
                                           loc=self._add(repr(value)))
+        elif value is numpy.int32:
+            typ = builtins.fn_int32()
+            return asttyped.NameConstantT(value=None, type=typ,
+                                          loc=self._add("numpy.int32"))
+        elif value is numpy.int64:
+            typ = builtins.fn_int64()
+            return asttyped.NameConstantT(value=None, type=typ,
+                                          loc=self._add("numpy.int64"))
+        elif value is numpy.array:
+            typ = builtins.fn_array()
+            return asttyped.NameConstantT(value=None, type=typ,
+                                          loc=self._add("numpy.array"))
+        elif value is numpy.full:
+            typ = builtins.fn_make_array()
+            return asttyped.NameConstantT(value=None, type=typ,
+                                          loc=self._add("numpy.full"))
         elif isinstance(value, (int, float)):
             if isinstance(value, int):
                 typ = builtins.TInt()
@@ -637,9 +653,13 @@ def __init__(self, core, dmgr, engine=None):
         self.name = ""
         self.typedtree = []
         self.inject_at = 0
+        self.globals = {}
+
+        # We don't want some things from the prelude as they are provided in
+        # the host Python namespace and gain special meaning when quoted.
         self.prelude = prelude.globals()
         self.prelude.pop("print")
-        self.globals = {}
+        self.prelude.pop("array")
 
         self.functions = {}
 

diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -1681,6 +1681,21 @@ def body_gen(index):
                 return self.append(ir.Select(cond, arg0, arg1))
             else:
                 assert False
+        elif types.is_builtin(typ, "make_array"):
+            if len(node.args) == 2 and len(node.keywords) == 0:
+                arg0, arg1 = map(self.visit, node.args)
+
+                result = self.append(ir.Alloc([arg0], node.type))
+                def body_gen(index):
+                    self.append(ir.SetElem(result, index, arg1))
+                    return self.append(ir.Arith(ast.Add(loc=None), index,
+                                                ir.Constant(1, arg0.type)))
+                self._make_loop(ir.Constant(0, self._size_type),
+                                lambda index: self.append(ir.Compare(ast.Lt(loc=None), index, arg0)),
+                                body_gen)
+                return result
+            else:
+                assert False
         elif types.is_builtin(typ, "print"):
             self.polymorphic_print([self.visit(arg) for arg in node.args],
                                    separator=" ", suffix="\n")
@@ -1725,7 +1740,7 @@ def body_gen(index):
         else:
             diag = diagnostic.Diagnostic("error",
                 "builtin function '{name}' cannot be used in this context",
-                {"name": typ.name},
+                {"name": typ.find().name},
                 node.loc)
             self.engine.process(diag)
 

diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py
@@ -845,6 +845,23 @@ def makenotes(printer, typea, typeb, loca, locb):
                 pass
             else:
                 diagnose(valid_forms())
+        elif types.is_builtin(typ, "make_array"):
+            valid_forms = lambda: [
+                valid_form("numpy.full(count:int32, value:'a) -> numpy.array(elt='a)")
+            ]
+
+            self._unify(node.type, builtins.TArray(),
+                        node.loc, None)
+
+            if len(node.args) == 2 and len(node.keywords) == 0:
+                arg0, arg1 = node.args
+
+                self._unify(arg0.type, builtins.TInt32(),
+                            arg0.loc, None)
+                self._unify(arg1.type, node.type.find()["elt"],
+                            arg1.loc, None)
+            else:
+                diagnose(valid_forms())
         elif types.is_builtin(typ, "rtio_log"):
             valid_forms = lambda: [
                 valid_form("rtio_log(channel:str, args...) -> None"),

diff --git a/artiq/test/coredevice/test_embedding.py b/artiq/test/coredevice/test_embedding.py
@@ -105,6 +105,14 @@ def kwargs2(self):
     def args1kwargs2(self):
         return self.kwargs("X", a="A", b=1)
 
+    @kernel
+    def numpy_things(self):
+        return (numpy.int32(10), numpy.int64(20), numpy.array([42,]))
+
+    @kernel
+    def numpy_full(self):
+        return numpy.full(10, 20)
+
     @kernel
     def builtin(self):
         sleep(1.0)
@@ -120,6 +128,9 @@ def test_args(self):
         self.assertEqual(exp.kwargs1(), 1)
         self.assertEqual(exp.kwargs2(), 2)
         self.assertEqual(exp.args1kwargs2(), 2)
+        self.assertEqual(exp.numpy_things(),
+                         (numpy.int32(10), numpy.int64(20), numpy.array([42,])))
+        self.assertTrue((exp.numpy_full() == numpy.full(10, 20)).all())
         exp.builtin()