From d882f8a3f031b337ad4008b4dfef48ea52cb258e Mon Sep 17 00:00:00 2001 From: David Nadlinger Date: Sat, 25 Jul 2020 22:34:52 +0100 Subject: [PATCH] compiler: Implement len() for ndarrays --- artiq/compiler/transforms/llvm_ir_generator.py | 4 ++++ artiq/test/lit/integration/array.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py index 3da6037a1..0ed3581fc 100644 --- a/artiq/compiler/transforms/llvm_ir_generator.py +++ b/artiq/compiler/transforms/llvm_ir_generator.py @@ -1170,6 +1170,10 @@ class LLVMIRGenerator: return get_outer(self.map(env), env.type) elif insn.op == "len": collection, = insn.operands + if builtins.is_array(collection.type): + # Return length of outermost dimension. + shape = self.llbuilder.extract_value(self.map(collection), 0) + return self.llbuilder.load(self.llbuilder.extract_value(shape, 0)) return self.llbuilder.extract_value(self.map(collection), 1) elif insn.op in ("printf", "rtio_log"): # We only get integers, floats, pointers and strings here. diff --git a/artiq/test/lit/integration/array.py b/artiq/test/lit/integration/array.py index 3cb7fc58d..c85ebaac1 100644 --- a/artiq/test/lit/integration/array.py +++ b/artiq/test/lit/integration/array.py @@ -2,7 +2,9 @@ # REQUIRES: exceptions ary = array([1, 2, 3]) +assert len(ary) == 3 # FIXME: Implement ndarray indexing # assert [x*x for x in ary] == [1, 4, 9] matrix = array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) +assert len(matrix) == 2