Implementing numpy.mod for array inputs

integration_tests/CMakeLists.txt

@@ -177,6 +177,7 @@ RUN(NAME elemental_06        LABELS cpython llvm)
 RUN(NAME elemental_07        LABELS cpython llvm)
 RUN(NAME elemental_08        LABELS cpython llvm)
 RUN(NAME elemental_09        LABELS cpython llvm)
+RUN(NAME elemental_10        LABELS cpython llvm)
 RUN(NAME elemental_11        LABELS cpython llvm)
 RUN(NAME test_random         LABELS cpython llvm)
 RUN(NAME test_os             LABELS cpython llvm)

integration_tests/elemental_10.py

@@ -0,0 +1,78 @@
+from ltypes import i32, i64
+from numpy import mod, int64, empty
+
+def test_numpy_mod():
+    q1: i64[32, 16, 7] = empty((32, 16, 7), dtype=int64)
+    d1: i64[32, 16, 7] = empty((32, 16, 7), dtype=int64)
+    r1: i64[32, 16, 7] = empty((32, 16, 7), dtype=int64)
+    r1neg: i64[32, 16, 7] = empty((32, 16, 7), dtype=int64)
+    q2: i64[100] = empty(100, dtype=int64)
+    d2: i64[100] = empty(100, dtype=int64)
+    r2: i64[100] = empty(100, dtype=int64)
+    i: i32; j: i32; k: i32
+    rem: i64; q: i64; d: i64
+
+    for i in range(32):
+        for j in range(16):
+            for k in range(7):
+                d1[i, j, k] = k + 1
+                q1[i, j, k] = (i + j) * (k + 1) + k
+
+    r1 = mod(q1, d1)
+    r1neg = mod(-q1, d1)
+
+    for i in range(32):
+        for j in range(16):
+            for k in range(7):
+                assert r1[i, j, k] == k
+                if k == 0:
+                    rem = 0
+                else:
+                    rem = d1[i, j, k] - k
+                assert r1neg[i, j, k] == rem
+
+    for i in range(32):
+        for j in range(16):
+            for k in range(7):
+                d1[i, j, k] = k + 2
+                q1[i, j, k] = i + j
+
+    r1 = mod(d1 * q1 + r1 + 1, d1)
+
+    for i in range(32):
+        for j in range(16):
+            for k in range(7):
+                assert r1[i, j, k] == k + 1
+
+    r1 = mod(2 * q1 + 1, int(2))
+
+    for i in range(32):
+        for j in range(16):
+            for k in range(7):
+                assert r1[i, j, k] == 1
+
+    for i in range(100):
+        d2[i] = i + 1
+
+    r2 = mod(int(100), d2)
+
+    for i in range(100):
+        assert r2[i] == 100 % (i + 1)
+
+    for i in range(100):
+        d2[i] = 50 - i
+        q2[i] = 39 - i
+
+    r2 = mod(q2, d2)
+
+    for i in range(100):
+        d = 50 - i
+        q = 39 - i
+        rem = r2[i]
+        if d == 0:
+            assert rem == 0
+        else:
+            assert int((q - rem)/d) - (q - rem)/d == 0
+
+
+test_numpy_mod()

src/libasr/pass/array_op.cpp

@@ -870,66 +870,99 @@ class ArrayOpVisitor : public PassUtils::PassVisitor<ArrayOpVisitor>
                                                 sub, nullptr,
                                                 s_args.p, s_args.size(), nullptr));
             pass_result.push_back(al, subrout_call);
-        } else if( is_elemental(x.m_name) && x.n_args == 1 &&
-                   ASRUtils::is_array(ASRUtils::expr_type(x.m_args[0].m_value)) ) {
+        } else if( is_elemental(x.m_name) ) {
+            std::vector<bool> array_mask(x.n_args, false);
+            bool at_least_one_array = false;
+            for( size_t iarg = 0; iarg < x.n_args; iarg++ ) {
+                array_mask[iarg] = ASRUtils::is_array(
+                    ASRUtils::expr_type(x.m_args[iarg].m_value));
+                at_least_one_array = at_least_one_array || array_mask[iarg];
+            }
+            if (!at_least_one_array) {
+                return ;
+            }
             std::string res_prefix = "_elemental_func_call_res";
             ASR::expr_t* result_var_copy = result_var;
-            result_var = nullptr;
-            this->visit_expr(*(x.m_args[0].m_value));
-            ASR::expr_t* operand = tmp_val;
-            int rank_operand = PassUtils::get_rank(operand);
-            if( rank_operand == 0 ) {
+            bool is_all_rank_0 = true;
+            std::vector<ASR::expr_t*> operands;
+            ASR::expr_t* operand = nullptr;
+            int common_rank = 0;
+            bool are_all_rank_same = true;
+            for( size_t iarg = 0; iarg < x.n_args; iarg++ ) {
+                result_var = nullptr;
+                this->visit_expr(*(x.m_args[iarg].m_value));
+                operand = tmp_val;
+                operands.push_back(operand);
+                int rank_operand = PassUtils::get_rank(operand);
+                if( common_rank == 0 ) {
+                    common_rank = rank_operand;
+                }
+                if( common_rank != rank_operand &&
+                    rank_operand > 0 ) {
+                    are_all_rank_same = false;
+                }
+                array_mask[iarg] = (rank_operand > 0);
+                is_all_rank_0 = is_all_rank_0 && (rank_operand <= 0);
+            }
+            if( is_all_rank_0 ) {
                 tmp_val = const_cast<ASR::expr_t*>(&(x.base));
                 return ;
             }
-            if( rank_operand > 0 ) {
-                result_var = result_var_copy;
-                if( result_var == nullptr ) {
-                    result_var = create_var(result_var_num, res_prefix,
-                                            x.base.base.loc, operand);
-                    result_var_num += 1;
-                }
-                tmp_val = result_var;
-
-                int n_dims = rank_operand;
-                Vec<ASR::expr_t*> idx_vars;
-                PassUtils::create_idx_vars(idx_vars, n_dims, x.base.base.loc, al, current_scope);
-                ASR::stmt_t* doloop = nullptr;
-                for( int i = n_dims - 1; i >= 0; i-- ) {
-                    // TODO: Add an If debug node to check if the lower and upper bounds of both the arrays are same.
-                    ASR::do_loop_head_t head;
-                    head.m_v = idx_vars[i];
-                    head.m_start = PassUtils::get_bound(result_var, i + 1, "lbound", al);
-                    head.m_end = PassUtils::get_bound(result_var, i + 1, "ubound", al);
-                    head.m_increment = nullptr;
-                    head.loc = head.m_v->base.loc;
-                    Vec<ASR::stmt_t*> doloop_body;
-                    doloop_body.reserve(al, 1);
-                    if( doloop == nullptr ) {
-                        ASR::expr_t* ref = PassUtils::create_array_ref(operand, idx_vars, al);
-                        ASR::expr_t* res = PassUtils::create_array_ref(result_var, idx_vars, al);
-                        ASR::expr_t* op_el_wise = nullptr;
+            if( !are_all_rank_same ) {
+                throw LCompilersException("Broadcasting support not yet available "
+                                          "for different shape arrays.");
+            }
+            result_var = result_var_copy;
+            if( result_var == nullptr ) {
+                result_var = create_var(result_var_num, res_prefix,
+                                        x.base.base.loc, operand);
+                result_var_num += 1;
+            }
+            tmp_val = result_var;
+
+            int n_dims = common_rank;
+            Vec<ASR::expr_t*> idx_vars;
+            PassUtils::create_idx_vars(idx_vars, n_dims, x.base.base.loc, al, current_scope);
+            ASR::stmt_t* doloop = nullptr;
+            for( int i = n_dims - 1; i >= 0; i-- ) {
+                // TODO: Add an If debug node to check if the lower and upper bounds of both the arrays are same.
+                ASR::do_loop_head_t head;
+                head.m_v = idx_vars[i];
+                head.m_start = PassUtils::get_bound(result_var, i + 1, "lbound", al);
+                head.m_end = PassUtils::get_bound(result_var, i + 1, "ubound", al);
+                head.m_increment = nullptr;
+                head.loc = head.m_v->base.loc;
+                Vec<ASR::stmt_t*> doloop_body;
+                doloop_body.reserve(al, 1);
+                if( doloop == nullptr ) {
+                    Vec<ASR::call_arg_t> ref_args;
+                    ref_args.reserve(al, x.n_args);
+                    for( size_t iarg = 0; iarg < x.n_args; iarg++ ) {
+                        ASR::expr_t* ref = operands[iarg];
+                        if( array_mask[iarg] ) {
+                            ref = PassUtils::create_array_ref(operands[iarg], idx_vars, al);
+                        }
                         ASR::call_arg_t ref_arg;
                         ref_arg.loc = ref->base.loc;
                         ref_arg.m_value = ref;
-                        Vec<ASR::call_arg_t> ref_args;
-                        ref_args.reserve(al, 1);
                         ref_args.push_back(al, ref_arg);
-                        Vec<ASR::dimension_t> empty_dim;
-                        empty_dim.reserve(al, 1);
-                        ASR::ttype_t* dim_less_type = ASRUtils::duplicate_type(al, x.m_type, &empty_dim);
-                        op_el_wise = LFortran::ASRUtils::EXPR(ASR::make_FunctionCall_t(al, x.base.base.loc,
-                                        x.m_name, x.m_original_name, ref_args.p, ref_args.size(), dim_less_type,
-                                        nullptr, x.m_dt));
-                        ASR::stmt_t* assign = LFortran::ASRUtils::STMT(ASR::make_Assignment_t(al, x.base.base.loc, res, op_el_wise, nullptr));
-                        doloop_body.push_back(al, assign);
-                    } else {
-                        doloop_body.push_back(al, doloop);
                     }
-                    doloop = LFortran::ASRUtils::STMT(ASR::make_DoLoop_t(al, x.base.base.loc, head, doloop_body.p, doloop_body.size()));
+                    Vec<ASR::dimension_t> empty_dim;
+                    empty_dim.reserve(al, 1);
+                    ASR::ttype_t* dim_less_type = ASRUtils::duplicate_type(al, x.m_type, &empty_dim);
+                    ASR::expr_t* op_el_wise = nullptr;
+                    op_el_wise = LFortran::ASRUtils::EXPR(ASR::make_FunctionCall_t(al, x.base.base.loc,
+                                    x.m_name, x.m_original_name, ref_args.p, ref_args.size(), dim_less_type,
+                                    nullptr, x.m_dt));
+                    ASR::expr_t* res = PassUtils::create_array_ref(result_var, idx_vars, al);
+                    ASR::stmt_t* assign = LFortran::ASRUtils::STMT(ASR::make_Assignment_t(al, x.base.base.loc, res, op_el_wise, nullptr));
+                    doloop_body.push_back(al, assign);
+                } else {
+                    doloop_body.push_back(al, doloop);
                 }
-                pass_result.push_back(al, doloop);
+                doloop = LFortran::ASRUtils::STMT(ASR::make_DoLoop_t(al, x.base.base.loc, head, doloop_body.p, doloop_body.size()));
             }
+            pass_result.push_back(al, doloop);
         }
         result_var = nullptr;
     }

src/libasr/pass/pass_utils.h

@@ -142,6 +142,10 @@ namespace LFortran {
                             ASR::AssociateBlock_t *s = ASR::down_cast<ASR::AssociateBlock_t>(item.second);
                             self().visit_AssociateBlock(*s);
                         }
+                        if (ASR::is_a<ASR::Block_t>(*item.second)) {
+                            ASR::Block_t *s = ASR::down_cast<ASR::Block_t>(item.second);
+                            self().visit_Block(*s);
+                        }
                     }
                 }
 
@@ -151,6 +155,13 @@ namespace LFortran {
                     ASR::Function_t &xx = const_cast<ASR::Function_t&>(x);
                     current_scope = xx.m_symtab;
                     transform_stmts(xx.m_body, xx.n_body);
+
+                    for (auto &item : x.m_symtab->get_scope()) {
+                        if (ASR::is_a<ASR::Block_t>(*item.second)) {
+                            ASR::Block_t *s = ASR::down_cast<ASR::Block_t>(item.second);
+                            self().visit_Block(*s);
+                        }
+                    }
                 }
 
                 void visit_AssociateBlock(const ASR::AssociateBlock_t& x) {
@@ -163,6 +174,10 @@ namespace LFortran {
                     ASR::Block_t &xx = const_cast<ASR::Block_t&>(x);
                     current_scope = xx.m_symtab;
                     transform_stmts(xx.m_body, xx.n_body);
+
+                    for (auto &item : x.m_symtab->get_scope()) {
+                        self().visit_symbol(*item.second);
+                    }
                 }
 
         };

src/lpython/semantics/python_ast_to_asr.cpp

@@ -862,21 +862,18 @@ class CommonVisitor : public AST::BaseVisitor<Derived> {
             }
 
             SymbolTable *symtab = current_scope;
-            while (symtab->parent != nullptr && symtab->get_scope().find(local_sym) == symtab->get_scope().end()) {
-                symtab = symtab->parent;
-            }
-            if (symtab->get_scope().find(local_sym) == symtab->get_scope().end()) {
+            if (symtab->resolve_symbol(local_sym) == nullptr) {
                 LFORTRAN_ASSERT(ASR::is_a<ASR::ExternalSymbol_t>(*stemp));
                 std::string mod_name = ASR::down_cast<ASR::ExternalSymbol_t>(stemp)->m_module_name;
-                ASR::symbol_t *mt = symtab->get_symbol(mod_name);
+                ASR::symbol_t *mt = symtab->resolve_symbol(mod_name);
                 ASR::Module_t *m = ASR::down_cast<ASR::Module_t>(mt);
                 stemp = import_from_module(al, m, symtab, mod_name,
                                     remote_sym, local_sym, loc);
                 LFORTRAN_ASSERT(ASR::is_a<ASR::ExternalSymbol_t>(*stemp));
                 symtab->add_symbol(local_sym, stemp);
                 s = ASRUtils::symbol_get_past_external(stemp);
             } else {
-                stemp = symtab->get_symbol(local_sym);
+                stemp = symtab->resolve_symbol(local_sym);
             }
         }
         if (ASR::is_a<ASR::Function_t>(*s)) {
@@ -3419,9 +3416,6 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
         current_scope = al.make_new<SymbolTable>(parent_scope);
         transform_stmts(body, x.n_body, x.m_body);
         int32_t total_syms = current_scope->get_scope().size();
-        for( auto& item: current_scope->get_scope() ) {
-            total_syms -= ASR::is_a<ASR::ExternalSymbol_t>(*item.second);
-        }
         if( total_syms > 0 ) {
             std::string name = parent_scope->get_unique_name("block");
             ASR::asr_t* block = ASR::make_Block_t(al, x.base.base.loc,
@@ -3433,24 +3427,8 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
                 ASR::down_cast<ASR::symbol_t>(block)));
             body.reserve(al, 1);
             body.push_back(al, decls);
-        } else {
-            // Revert global counter as no variables
-            // were declared inside the loop so
-            // current_scope is not needed.
-            for( auto& item: current_scope->get_scope() ) {
-                if( !ASR::is_a<ASR::ExternalSymbol_t>(*item.second) ) {
-                    continue ;
-                }
-
-                ASR::ExternalSymbol_t* ext_sym = ASR::down_cast<ASR::ExternalSymbol_t>(item.second);
-                ASR::symbol_t* new_ext_sym = ASR::down_cast<ASR::symbol_t>(
-                    ASR::make_ExternalSymbol_t(al, ext_sym->base.base.loc, parent_scope, ext_sym->m_name,
-                            ext_sym->m_external, ext_sym->m_module_name, ext_sym->m_scope_names,
-                            ext_sym->n_scope_names, ext_sym->m_original_name, ext_sym->m_access));
-                parent_scope->add_symbol(item.first, new_ext_sym);
-            }
-            current_scope = parent_scope;
         }
+        current_scope = parent_scope;
 
         if (loop_start) {
             head.m_start = loop_start;

src/runtime/lpython_intrinsic_numpy.py

@@ -1,4 +1,4 @@
-from ltypes import f64, f32, ccall, vectorize, overload
+from ltypes import i32, i64, f64, f32, ccall, vectorize, overload
 
 pi_64: f64 = 3.141592653589793238462643383279502884197
 pi_32: f32 = 3.141592653589793238462643383279502884197
@@ -354,3 +354,19 @@ def _lfortran_satanh(x: f32) -> f32:
 @vectorize
 def arctanh(x: f32) -> f32:
     return _lfortran_satanh(x)
+
+########## mod ##########
+
+@overload
+@vectorize
+def mod(x1: i64, x2: i64) -> i64:
+    if x2 == 0:
+        return int(0)
+    return x1 % x2
+
+@overload
+@vectorize
+def mod(x1: i32, x2: i32) -> i32:
+    if x2 == 0:
+        return 0
+    return x1 % x2

tests/reference/asr-array_01_decl-39cf894.json

@@ -6,7 +6,7 @@
     "outfile": null,
     "outfile_hash": null,
     "stdout": "asr-array_01_decl-39cf894.stdout",
-    "stdout_hash": "626eb3f1b687f1ba8499cb42d852d4e2fc316a11b2f48a058380c23a",
+    "stdout_hash": "2380e43101fe35534821c0b5707103f481413bb34c98d6c9ea7cf3dc",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0

tests/reference/asr-array_02_decl-e8f6874.json

@@ -6,7 +6,7 @@
     "outfile": null,
     "outfile_hash": null,
     "stdout": "asr-array_02_decl-e8f6874.stdout",
-    "stdout_hash": "c25f1a5906a7ff591373fdd8a67d5e879f359390030a27d9bfc3ee9a",
+    "stdout_hash": "dfcfd217678648da9009018c17a45aab3334b7d0625a80318066c2d7",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0

tests/reference/asr-complex1-f26c460.json

@@ -6,7 +6,7 @@
     "outfile": null,
     "outfile_hash": null,
     "stdout": "asr-complex1-f26c460.stdout",
-    "stdout_hash": "c741bb3cc0e4ba3e9076c989424b8eddb5ebdb9458b9fb894dc84044",
+    "stdout_hash": "c6f072f66291b41dec2703ad724fcacb63378c280ee00018739ec8b7",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0