Skip to content

Commit ff38f75

Browse files
certikcertik
authored
Merge pull request #1777 from Shaikh-Ubaid/numpy_size
Support numpy size() and array.size
2 parents 9b8fa1e + c900f7d commit ff38f75

8 files changed

+74
-24
lines changed

‎integration_tests/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,6 +280,7 @@ RUN(NAME variable_decl_02 LABELS cpython llvm c)
280280
RUN(NAME variable_decl_03 LABELS cpython llvm c)
281281
RUN(NAME array_expr_01 LABELS cpython llvm c)
282282
RUN(NAME array_expr_02 LABELS cpython llvm c)
283+
RUN(NAME array_size_01 LABELS cpython llvm c)
283284
RUN(NAME array_01 LABELS cpython llvm wasm c)
284285
RUN(NAME array_02 LABELS cpython wasm c)
285286
RUN(NAME bindc_01 LABELS cpython llvm c)

‎integration_tests/array_size_01.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
from lpython import i32, f64, c32, c64
2+
from numpy import empty
3+
4+
def main0():
5+
x: i32[4, 5, 2] = empty([4, 5, 2])
6+
y: f64[24, 100, 2, 5] = empty([24, 100, 2, 5])
7+
print(x.size)
8+
print(y.size)
9+
10+
assert x.size == 40
11+
assert y.size == 24000
12+
13+
def main1():
14+
a: c32[12] = empty([12])
15+
b: c64[15, 15, 10] = empty([15, 15, 10])
16+
print(a.size)
17+
print(b.size)
18+
19+
assert a.size == 12
20+
assert b.size == 2250
21+
22+
main0()
23+
main1()

‎src/lpython/semantics/python_ast_to_asr.cpp

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5137,7 +5137,13 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
51375137

51385138
void visit_AttributeUtil(ASR::ttype_t* type, char* attr_char,
51395139
ASR::symbol_t *t, const Location& loc) {
5140-
if (ASRUtils::is_complex(*type)) {
5140+
if (ASRUtils::is_array(type)) {
5141+
std::string attr = attr_char;
5142+
ASR::expr_t *se = ASR::down_cast<ASR::expr_t>(ASR::make_Var_t(al, loc, t));
5143+
Vec<ASR::expr_t*> args;
5144+
args.reserve(al, 0);
5145+
handle_attribute(se, attr, loc, args);
5146+
} else if (ASRUtils::is_complex(*type)) {
51415147
std::string attr = attr_char;
51425148
if (attr == "imag") {
51435149
ASR::expr_t *val = ASR::down_cast<ASR::expr_t>(ASR::make_Var_t(al, loc, t));
@@ -6705,7 +6711,7 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
67056711
ASRUtils::IntrinsicFunctionRegistry::get_create_function(call_name);
67066712
Vec<ASR::expr_t*> args_; args_.reserve(al, x.n_args);
67076713
visit_expr_list(x.m_args, x.n_args, args_);
6708-
if (ASRUtils::is_array(ASRUtils::expr_type(args_[0])) &&
6714+
if (ASRUtils::is_array(ASRUtils::expr_type(args_[0])) &&
67096715
imported_functions[call_name] == "math" ) {
67106716
throw SemanticError("Function '" + call_name + "' does not accept vector values",
67116717
x.base.base.loc);
@@ -6737,21 +6743,19 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
67376743
// This will all be removed once we port it to intrinsic functions
67386744
// Intrinsic functions
67396745
if (call_name == "size") {
6740-
// TODO: size should be part of ASR. That way
6741-
// ASR itself does not need a runtime library
6742-
// a runtime library thus becomes optional --- can either be
6743-
// implemented using ASR, or the backend can link it at runtime
6744-
ASR::ttype_t *a_type = ASRUtils::TYPE(ASR::make_Integer_t(al, x.base.base.loc,
6745-
4, nullptr, 0));
6746-
/*
6747-
ASR::symbol_t *a_name = nullptr;
6748-
throw SemanticError("TODO: add the size() function and look it up",
6749-
x.base.base.loc);
6750-
tmp = ASR::make_FunctionCall_t(al, x.base.base.loc, a_name,
6751-
nullptr, args.p, args.size(), nullptr, 0, a_type, nullptr, nullptr);
6752-
*/
6753-
6754-
tmp = ASR::make_IntegerConstant_t(al, x.base.base.loc, 1234, a_type);
6746+
parse_args();
6747+
if( args.size() < 1 || args.size() > 2 ) {
6748+
throw SemanticError("array accepts only 1 (arr) or 2 (arr, axis) arguments, got " +
6749+
std::to_string(args.size()) + " arguments instead.",
6750+
x.base.base.loc);
6751+
}
6752+
ASR::expr_t *var = args[0].m_value;
6753+
ASR::expr_t *dim = nullptr;
6754+
if (args.size() == 2) {
6755+
dim = args[1].m_value;
6756+
}
6757+
ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, x.base.base.loc, 4, nullptr, 0));
6758+
tmp = ASR::make_ArraySize_t(al, x.base.base.loc, var, dim, int_type, nullptr);
67556759
return;
67566760
} else if (call_name == "empty") {
67576761
// TODO: check that the `empty` arguments are compatible

‎src/lpython/semantics/python_attribute_eval.h

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ struct AttributeHandler {
2020
AttributeHandler() {
2121
attribute_map = {
2222
{"int@bit_length", &eval_int_bit_length},
23+
{"array@size", &eval_array_size},
2324
{"list@append", &eval_list_append},
2425
{"list@remove", &eval_list_remove},
2526
{"list@count", &eval_list_count},
@@ -42,6 +43,8 @@ struct AttributeHandler {
4243
return "set";
4344
} else if (ASR::is_a<ASR::Dict_t>(*t)) {
4445
return "dict";
46+
} else if (ASRUtils::is_array(t)) {
47+
return "array";
4548
} else if (ASR::is_a<ASR::Integer_t>(*t)) {
4649
return "int";
4750
}
@@ -55,7 +58,7 @@ struct AttributeHandler {
5558
if (class_name == "") {
5659
throw SemanticError("Type name is not implemented yet.", loc);
5760
}
58-
std::string key = get_type_name(type) + "@" + attr_name;
61+
std::string key = class_name + "@" + attr_name;
5962
auto search = attribute_map.find(key);
6063
if (search != attribute_map.end()) {
6164
attribute_eval_callback cb = search->second;
@@ -77,6 +80,15 @@ struct AttributeHandler {
7780
return ASR::make_IntegerBitLen_t(al, loc, s, int_type, nullptr);
7881
}
7982

83+
static ASR::asr_t* eval_array_size(ASR::expr_t *s, Allocator &al, const Location &loc,
84+
Vec<ASR::expr_t*> &args, diag::Diagnostics &/*diag*/) {
85+
if (args.size() != 0) {
86+
throw SemanticError("array.size() takes no arguments", loc);
87+
}
88+
ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, loc, 4, nullptr, 0));
89+
return ASR::make_ArraySize_t(al, loc, s, nullptr, int_type, nullptr);
90+
}
91+
8092
static ASR::asr_t* eval_list_append(ASR::expr_t *s, Allocator &al, const Location &loc,
8193
Vec<ASR::expr_t*> &args, diag::Diagnostics &diag) {
8294
if (args.size() != 1) {

‎tests/reference/asr-doconcurrentloop_01-7b9a7d3.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
"outfile": null,
77
"outfile_hash": null,
88
"stdout": "asr-doconcurrentloop_01-7b9a7d3.stdout",
9-
"stdout_hash": "dffb1f7233f48b1711345ac6d1649bfdc6d8244f5596e9a2a9144904",
9+
"stdout_hash": "9519c801f6f3612439fef115f1c30385234060565a7dd07e125685a4",
1010
"stderr": null,
1111
"stderr_hash": null,
1212
"returncode": 0

‎tests/reference/asr-doconcurrentloop_01-7b9a7d3.stdout

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,12 @@
179179
)
180180
(=
181181
(Var 3 nsize)
182-
(IntegerConstant 1234 (Integer 4 []))
182+
(ArraySize
183+
(Var 3 a)
184+
()
185+
(Integer 4 [])
186+
()
187+
)
183188
()
184189
)
185190
(DoConcurrentLoop
@@ -394,7 +399,12 @@
394399
(Var 2 c)]
395400
[(=
396401
(Var 2 N)
397-
(IntegerConstant 1234 (Integer 4 []))
402+
(ArraySize
403+
(Var 2 a)
404+
()
405+
(Integer 4 [])
406+
()
407+
)
398408
()
399409
)
400410
(DoConcurrentLoop

‎tests/reference/cpp-doconcurrentloop_01-4e9f274.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
"outfile": null,
77
"outfile_hash": null,
88
"stdout": "cpp-doconcurrentloop_01-4e9f274.stdout",
9-
"stdout_hash": "ee58f996458bdf91f906f012b009681a568ad314c5907d9a698e41b8",
9+
"stdout_hash": "79dc6e24e5f0c768ebbadd5b3bd466b53e1b89608b2d1e1a4a0d4caf",
1010
"stderr": null,
1111
"stderr_hash": null,
1212
"returncode": 0

‎tests/reference/cpp-doconcurrentloop_01-4e9f274.stdout

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ void triad(T0* a, T1* b, float scalar, T2* c)
4848
{
4949
int32_t N;
5050
int32_t i;
51-
N = 1234;
51+
N = a->data->extent(0);
5252
Kokkos::parallel_for(Kokkos::RangePolicy<Kokkos::DefaultExecutionSpace>(0, N - 1+1), KOKKOS_LAMBDA(const long i) {
5353
c->data->operator[](i - c->dims[0].lower_bound) = a->data->operator[](i - a->dims[0].lower_bound) + scalar*b->data->operator[](i - b->dims[0].lower_bound);
5454
});
@@ -75,7 +75,7 @@ void main0()
7575
int32_t nsize;
7676
float scalar;
7777
scalar = 1.00000000000000000e+01;
78-
nsize = 1234;
78+
nsize = a->data->extent(0);
7979
Kokkos::parallel_for(Kokkos::RangePolicy<Kokkos::DefaultExecutionSpace>(0, nsize - 1+1), KOKKOS_LAMBDA(const long i) {
8080
a->data->operator[](i - a->dims[0].lower_bound) = 5.00000000000000000e+00;
8181
b->data->operator[](i - b->dims[0].lower_bound) = 5.00000000000000000e+00;

0 commit comments

Comments
 (0)