compiler: Prioritize large sparse loops over tiny ones

devito/passes/iet/parpragma.py

@@ -266,23 +266,33 @@ def _score_candidate(self, n0, root, collapsable=()):
 
         # Number of fully-parallel collapsable Iterations
         key = lambda i: i.is_ParallelNoAtomic
-        fpiters = list(takewhile(key, nest))
-        nfpiters = len(fpiters)
+        fp_iters = list(takewhile(key, nest))
+        n_fp_iters = len(fp_iters)
+
+        # Number of parallel-if-atomic collapsable Iterations
+        key = lambda i: i.is_ParallelAtomic
+        pia_iters = list(takewhile(key, nest))
+        n_pia_iters = len(pia_iters)
 
         # Prioritize the Dimensions that are more likely to define larger
         # iteration spaces
-        fpdims = [i.dim for i in fpiters]
         key = lambda d: (not d.is_Derived or
-                         d.is_Custom or  # NOTE: might use a refinement
+                         (d.is_Custom and not is_integer(d.symbolic_size)) or
                          (d.is_Block and d._depth == 1))
-        nfpiters_large = len([d for d in fpdims if key(d)])
+
+        fpdims = [i.dim for i in fp_iters]
+        n_fp_iters_large = len([d for d in fpdims if key(d)])
+
+        piadims = [i.dim for i in pia_iters]
+        n_pia_iters_large = len([d for d in piadims if key(d)])
 
         return (
-            int(nfpiters == n),  # Fully-parallel nest
-            int(nfpiters == 0 and n),  # Fully-atomic nest
-            nfpiters_large,
+            int(n_fp_iters == n),  # Fully-parallel nest
+            n_fp_iters_large,
+            n_pia_iters_large,
+            n_pia_iters,
             -(n0 + 1),  # The outer, the better
-            nfpiters,
+            n_fp_iters,
             n,
         )
 

tests/test_dle.py

@@ -284,7 +284,7 @@ def test_cache_blocking_structure_optrelax_prec_inject():
                                           'openmp': True,
                                           'par-collapse-ncores': 1}))
 
-    assert_structure(op, ['t,p_s0_blk0,p_s', 't,p_s0_blk0,p_s,rsx,rsy'],
+    assert_structure(op, ['t', 't,p_s0_blk0,p_s,rsx,rsy'],
                      't,p_s0_blk0,p_s,rsx,rsy')
 
 
@@ -897,6 +897,23 @@ def test_incr_perfect_outer(self):
         op()
         assert np.all(w.data == 10)
 
+    def test_incr_perfect_sparse_outer(self):
+        grid = Grid(shape=(3, 3, 3))
+
+        u = TimeFunction(name='u', grid=grid)
+        s = SparseTimeFunction(name='u', grid=grid, npoint=1, nt=11)
+
+        eqns = s.inject(u, expr=s)
+
+        op = Operator(eqns, opt=('advanced', {'par-collapse-ncores': 0}))
+
+        iters = FindNodes(Iteration).visit(op)
+        assert len(iters) == 5
+        assert iters[0].is_Sequential
+        assert all(i.is_ParallelAtomic for i in iters[1:])
+        assert iters[1].pragmas[0].value == 'omp for schedule(dynamic,chunk_size)'
+        assert all(not i.pragmas for i in iters[2:])
+
     @pytest.mark.parametrize('exprs,simd_level,expected', [
         (['Eq(y.symbolic_max, g[0, x], implicit_dims=(t, x))',
          'Inc(h1[0, 0], 1, implicit_dims=(t, x, y))'],