Merge pull request #2600 from devitocodes/blocking

compiler: Blocked SubDimensions now use root dimension name

Author: FabioLuporini

devito/passes/clusters/blocking.py

@@ -354,15 +354,15 @@ def _derive_block_dims(self, clusters, prefix, d, blk_size_gen):
             except KeyError:
                 pass
 
-        base = self.sregistry.make_name(prefix=d.name)
+        base = self.sregistry.make_name(prefix=d.root.name)
 
-        name = self.sregistry.make_name(prefix="%s_blk" % base)
+        name = self.sregistry.make_name(prefix=f"{base}_blk")
         bd = BlockDimension(name, d, d.symbolic_min, d.symbolic_max, step)
         step = bd.step
         block_dims = [bd]
 
         for _ in range(1, self.levels):
-            name = self.sregistry.make_name(prefix="%s_blk" % base)
+            name = self.sregistry.make_name(prefix=f"{base}_blk")
             bd = BlockDimension(name, bd, bd, bd + bd.step - 1, size=step)
             block_dims.append(bd)
 

tests/test_dimension.py

@@ -2044,21 +2044,21 @@ def test_topofusion_w_subdims_conddims(self):
 
         # Check generated code -- expect the gsave equation to be scheduled together
         # in the same loop nest with the fsave equation
-        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'ix0_blk0'})
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'x2_blk0'})
         exprs = FindNodes(Expression).visit(bns['x0_blk0'])
         assert len(exprs) == 2
         assert exprs[0].write is f
         assert exprs[1].write is g
 
         exprs = FindNodes(Expression).visit(bns['x1_blk0'])
+        assert len(exprs) == 1
+        assert exprs[0].write is h
+
+        exprs = FindNodes(Expression).visit(bns['x2_blk0'])
         assert len(exprs) == 2
         assert exprs[0].write is fsave
         assert exprs[1].write is gsave
 
-        exprs = FindNodes(Expression).visit(bns['ix0_blk0'])
-        assert len(exprs) == 1
-        assert exprs[0].write is h
-
     def test_topofusion_w_subdims_conddims_v2(self):
         """
         Like `test_topofusion_w_subdims_conddims` but with more SubDomains,
@@ -2085,9 +2085,9 @@ def test_topofusion_w_subdims_conddims_v2(self):
 
         # Check generated code -- expect the gsave equation to be scheduled together
         # in the same loop nest with the fsave equation
-        bns, _ = assert_blocking(op, {'ix0_blk0', 'x0_blk0'})
-        assert len(FindNodes(Expression).visit(bns['ix0_blk0'])) == 3
-        exprs = FindNodes(Expression).visit(bns['x0_blk0'])
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0'})
+        assert len(FindNodes(Expression).visit(bns['x0_blk0'])) == 3
+        exprs = FindNodes(Expression).visit(bns['x1_blk0'])
         assert len(exprs) == 2
         assert exprs[0].write is fsave
         assert exprs[1].write is gsave
@@ -2118,19 +2118,19 @@ def test_topofusion_w_subdims_conddims_v3(self):
 
         # Check generated code -- expect the gsave equation to be scheduled together
         # in the same loop nest with the fsave equation
-        bns, _ = assert_blocking(op, {'ix0_blk0', 'x0_blk0', 'ix1_blk0'})
-        exprs = FindNodes(Expression).visit(bns['ix0_blk0'])
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'x2_blk0'})
+        exprs = FindNodes(Expression).visit(bns['x0_blk0'])
         assert len(exprs) == 2
         assert exprs[0].write is f
         assert exprs[1].write is g
 
-        exprs = FindNodes(Expression).visit(bns['x0_blk0'])
+        exprs = FindNodes(Expression).visit(bns['x2_blk0'])
         assert len(exprs) == 2
         assert exprs[0].write is fsave
         assert exprs[1].write is gsave
 
         # Additional nest due to anti-dependence
-        exprs = FindNodes(Expression).visit(bns['ix1_blk0'])
+        exprs = FindNodes(Expression).visit(bns['x1_blk0'])
         assert len(exprs) == 2
         assert exprs[1].write is h
 

tests/test_dle.py

@@ -130,9 +130,9 @@ def test_cache_blocking_structure_subdims():
     # Non-local SubDimension -> blocking expected
     op = Operator(Eq(f.forward, f.dx + 1, subdomain=grid.interior))
 
-    bns, _ = assert_blocking(op, {'ix0_blk0'})
+    bns, _ = assert_blocking(op, {'x0_blk0'})
 
-    trees = retrieve_iteration_tree(bns['ix0_blk0'])
+    trees = retrieve_iteration_tree(bns['x0_blk0'])
     tree = trees[0]
     assert len(tree) == 5
     assert tree[0].dim.is_Block and tree[0].dim.parent.name == 'ix' and\
@@ -257,7 +257,7 @@ def test_leftright_subdims(self):
 
         op = Operator(eqns, opt=('fission', 'blocking', {'blockrelax': 'device-aware'}))
 
-        bns, _ = assert_blocking(op, {'x0_blk0', 'xl0_blk0', 'xr0_blk0'})
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'x2_blk0'})
         assert all(IsPerfectIteration().visit(i) for i in bns.values())
         assert all(len(FindNodes(Iteration).visit(i)) == 4 for i in bns.values())
 
@@ -1362,9 +1362,9 @@ def test_nested_cache_blocking_structure_subdims(self, blocklevels):
                             'par-collapse-ncores': 2,
                             'par-dynamic-work': 0}))
 
-        bns, _ = assert_blocking(op, {'ix0_blk0'})
+        bns, _ = assert_blocking(op, {'x0_blk0'})
 
-        trees = retrieve_iteration_tree(bns['ix0_blk0'])
+        trees = retrieve_iteration_tree(bns['x0_blk0'])
         assert len(trees) == 1
         tree = trees[0]
         assert len(tree) == 5 + (blocklevels - 1) * 2

tests/test_dse.py

@@ -599,11 +599,11 @@ def test_full_shape_w_subdims(self, rotate):
                                               'cire-rotate': rotate}))
 
         # Check code generation
-        bns, pbs = assert_blocking(op1, {'ix0_blk0'})
-        xs, ys, zs = get_params(op1, 'ix0_blk0_size', 'iy0_blk0_size', 'z_size')
-        arrays = [i for i in FindSymbols().visit(bns['ix0_blk0']) if i.is_Array]
+        bns, pbs = assert_blocking(op1, {'x0_blk0'})
+        xs, ys, zs = get_params(op1, 'x0_blk0_size', 'y0_blk0_size', 'z_size')
+        arrays = [i for i in FindSymbols().visit(bns['x0_blk0']) if i.is_Array]
         assert len(arrays) == 1
-        assert len(FindNodes(VExpanded).visit(pbs['ix0_blk0'])) == 1
+        assert len(FindNodes(VExpanded).visit(pbs['x0_blk0'])) == 1
         check_array(arrays[0], ((1, 1), (1, 1), (1, 1)), (xs+2, ys+2, zs+2), rotate)
 
         # Check numerical output
@@ -773,11 +773,11 @@ def test_mixed_shapes_v2_w_subdims(self, rotate):
                                   'cire-mingain': 0, 'cire-rotate': rotate}))
 
         # Check code generation
-        bns, pbs = assert_blocking(op1, {'ix0_blk0'})
-        xs, ys, zs = get_params(op1, 'ix0_blk0_size', 'iy0_blk0_size', 'z_size')
-        arrays = [i for i in FindSymbols().visit(bns['ix0_blk0']) if i.is_Array]
+        bns, pbs = assert_blocking(op1, {'x0_blk0'})
+        xs, ys, zs = get_params(op1, 'x0_blk0_size', 'y0_blk0_size', 'z_size')
+        arrays = [i for i in FindSymbols().visit(bns['x0_blk0']) if i.is_Array]
         assert len(arrays) == 2
-        assert len(FindNodes(VExpanded).visit(pbs['ix0_blk0'])) == 2
+        assert len(FindNodes(VExpanded).visit(pbs['x0_blk0'])) == 2
         check_array(arrays[0], ((1, 0), (1, 0), (0, 0)), (xs+1, ys+1, zs), rotate)
         check_array(arrays[1], ((1, 1), (1, 0)), (ys+2, zs+1), rotate)
 

tests/test_mpi.py

@@ -1279,11 +1279,11 @@ def test_hoist_haloupdate_with_subdims(self, mode):
         assert calls[1].name == 'haloupdate0'
 
         # ... and none in the created efuncs
-        bns, _ = assert_blocking(op, {'ix0_blk0', 'x0_blk0'})
-        calls = FindNodes(Call).visit(bns['ix0_blk0'])
-        assert len(calls) == 0
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0'})
         calls = FindNodes(Call).visit(bns['x0_blk0'])
         assert len(calls) == 0
+        calls = FindNodes(Call).visit(bns['x1_blk0'])
+        assert len(calls) == 0
 
     @pytest.mark.parallel(mode=1)
     def test_hoist_haloupdate_from_innerloop(self, mode):

tests/test_subdomains.py

@@ -680,8 +680,8 @@ class Dummy(SubDomainSet):
         # Make sure it jit-compiles
         op.cfunction
 
-        assert_structure(op, ['t,n0', 't,n0,ix0_blk0,iy0_blk0,x,y,z'],
-                         't,n0,ix0_blk0,iy0_blk0,x,y,z')
+        assert_structure(op, ['t,n0', 't,n0,x0_blk0,y0_blk0,x,y,z'],
+                         't,n0,x0_blk0,y0_blk0,x,y,z')
 
         # Drag a rebuilt MultiSubDimension out of the operator
         dims = {d.name: d for d in FindSymbols('dimensions').visit(op)}