import cloudpickle as pickle import pytest from itertools import product import numpy as np import sympy from tests.common import (get_cuda_graphs, init_chessboard, opt0, opt1, opt2, opt3, opt4, opt5, opt6, opt7, optX) from tests.conftest import skipif from devito import ( Buffer, ConditionalDimension, Constant, Dimension, Function, Grid, DefaultDimension, SparseFunction, SparseTimeFunction, PrecomputedSparseFunction, TimeFunction, TensorFunction, VectorFunction, VectorTimeFunction, TensorTimeFunction, CustomDimension, SubDomain, SubDimension, Eq, Inc, Operator, cos, sin, norm, switchconfig, configuration, div, diag, grad, NODE, exp, solve, Real, Imag, Conj, Derivative, Gt, Ge, Lt, Le, Ne, CondEq ) from devito.arch import AmdDevice from devito.data import CENTER, LEFT, RIGHT from devito.exceptions import ( # noqa CompilationError, InvalidArgument, InvalidOperator, ExecutionError ) from devito.finite_differences import Coefficient, Substitutions from devito.finite_differences.differentiable import SafeInv, diff2sympy from devito.ir import (Call, Conditional, DummyExpr, Expression, Iteration, Pragma, KernelLaunch, FindNodes, FindSymbols, SymbolRegistry, TimedList, retrieve_iteration_tree) from devito.mpi.routines import ( HaloUpdate, HaloUpdateCall, HaloUpdateList, HaloWait, HaloWaitList, MPIMsg, RemainderCall ) from devito.symbolics import (Cast, INT, ListInitializer, retrieve_derivatives, normalize_args) from devito.tools import as_tuple from devito.types import (DeviceID, DeviceRM, ModuloDimension, StencilDimension, Symbol, TBArray, TempArray) from devito.types.dimension import Thickness from devitopro import * # noqa from devitopro.passes import kernelopt from devitopro.passes.flip import _flip from devitopro.passes.languages import cuda from devitopro.passes.execution_graph import CudaGraphExec, CudaGraphKernelCall from devitopro.passes.parlang import ( ParlangTransformer, Stream, block_idxs, derive_grid_shape, thread_idxs ) from devitopro.passes.rotation_gather import QueueInitDimension from devitopro.types.array import TBQueue, TBShm from devitopro.types.dimension import (ThreadDimension, ThreadDimLayoutA, ThreadDimLayoutB, WarpId) from devitopro.types.enriched import DiskHost, NoLayers from examples.seismic import Model, RickerSource, TimeAxis pytestmark = skipif(['nodevice-lang'], whole_module=True) class TestBasic: @pytest.mark.parametrize('sizes,shape,indices', [ ((8, 24), (3,), (0, block_idxs[0])), ((24, 8, 16), (3, 2), (block_idxs[0], 0, block_idxs[1])), ((8, 8), (1,), (0, 0)), ]) def test_grid_shape(self, sizes, shape, indices): # Include nonzero origins when identifying single-block axes iters = [Iteration((), Dimension(f'd{i}'), (3, size + 2, 8)) for i, size in enumerate(sizes)] actual, subs = derive_grid_shape(iters, compact=True) assert sympy.lambdify((), actual, {'ceil': np.ceil})() == shape assert tuple(subs.get(i, i) for i in block_idxs[:len(sizes)]) == indices # Keep all axes without compaction actual, subs = derive_grid_shape(iters) expected = tuple(s//8 for s in sizes) assert sympy.lambdify((), actual, {'ceil': np.ceil})() == expected assert not subs @pytest.mark.parametrize('optoptions,expected', [ ({'cse-algo': 'advanced', 'fact-schedule': 'basic'}, {'cse-algo': ['advanced', 'advanced'], 'fact-schedule': ['basic', 'basic']}), ({'cse-algo': ['basic', 'advanced'], 'fact-schedule': ('smart', 'basic')}, {'cse-algo': ['basic', 'advanced'], 'fact-schedule': ['smart', 'basic']}), ], ids=['scalar', 'local']) def test_local_flops_options(self, monkeypatch, optoptions, expected): grid = Grid(shape=(8, 8, 8)) u = TimeFunction(name='u', grid=grid, space_order=2) v = TimeFunction(name='v', grid=grid, space_order=2) eqns = [Eq(u.forward, u + v.dy + 1), Eq(v.forward, u.forward.dx + v.dy + 1)] seen = {'cse-algo': [], 'fact-schedule': []} factorize = kernelopt.factorize cse = kernelopt.cse def track_factorize(clusters, options=None, **kwargs): if isinstance(clusters, kernelopt.Kernel): value = options['fact-schedule'] assert isinstance(value, str) seen['fact-schedule'].append(value) return factorize(clusters, options=options, **kwargs) def track_cse(clusters, options=None, **kwargs): if isinstance(clusters, kernelopt.Kernel): value = options['cse-algo'] assert isinstance(value, str) seen['cse-algo'].append(value) return cse(clusters, options=options, **kwargs) monkeypatch.setattr(kernelopt, 'factorize', track_factorize) monkeypatch.setattr(kernelopt, 'cse', track_cse) op = Operator(eqns, opt=opt0(optoptions)) assert seen == expected u.data_with_halo[:] = 2 v.data_with_halo[:] = 3 op.apply(time_M=0) assert np.allclose(u.data[1, 1:-1, 1:-1, 1:-1], 3) assert np.allclose(v.data[1, 1:-1, 1:-1, 1:-1], 1) @pytest.mark.parametrize('optoptions,expected', [ ({'deriv-uniformize': True, 'deriv-schedule': 'smart', 'deriv-unroll': 'full'}, {'deriv-uniformize': [True, True], 'deriv-schedule': ['smart', 'smart'], 'deriv-unroll': ['full', 'full']}), ({'deriv-uniformize': [False, True], 'deriv-schedule': ['basic', 'smart'], 'deriv-unroll': (False, 'full')}, {'deriv-uniformize': [False, True], 'deriv-schedule': ['basic', 'smart'], 'deriv-unroll': [False, 'full']}), ], ids=['scalar', 'local']) def test_local_derivative_options(self, monkeypatch, optoptions, expected): grid = Grid(shape=(8, 8, 8)) u = TimeFunction(name='u', grid=grid, space_order=2) v = TimeFunction(name='v', grid=grid, space_order=2) eqns = [Eq(u.forward, u + v.dy + 1), Eq(v.forward, u.forward.dx + v.dy + 1)] seen = {'deriv-uniformize': [], 'deriv-schedule': [], 'deriv-unroll': []} uniformize_index_derivatives = kernelopt.uniformize_index_derivatives specialize_index_derivatives = kernelopt.specialize_index_derivatives def track_uniformize(clusters, options=None, **kwargs): if isinstance(clusters, kernelopt.Kernel): uniformize = options['deriv-uniformize'] assert isinstance(uniformize, bool) seen['deriv-uniformize'].append(uniformize) return uniformize_index_derivatives(clusters, options=options, **kwargs) def track_specialize(clusters, options=None, **kwargs): if isinstance(clusters, kernelopt.Kernel): schedule = options['deriv-schedule'] unroll = options['deriv-unroll'] assert isinstance(schedule, str) assert isinstance(unroll, (bool, str)) seen['deriv-schedule'].append(schedule) seen['deriv-unroll'].append(unroll) return specialize_index_derivatives(clusters, options=options, **kwargs) monkeypatch.setattr(kernelopt, 'uniformize_index_derivatives', track_uniformize) monkeypatch.setattr(kernelopt, 'specialize_index_derivatives', track_specialize) op = Operator(eqns, opt=opt0(optoptions)) assert seen == expected u.data_with_halo[:] = 2 v.data_with_halo[:] = 3 op.apply(time_M=0) assert np.allclose(u.data[1, 1:-1, 1:-1, 1:-1], 3) assert np.allclose(v.data[1, 1:-1, 1:-1, 1:-1], 1) @pytest.mark.parametrize('name,value', [ ('cse-algo', ('basic', 'illegal')), ('fact-schedule', ('smart', 'illegal')), ('deriv-uniformize', (False, 'illegal')), ('deriv-schedule', ('smart', 'illegal')), ('deriv-unroll', ('full', 'illegal')), ]) def test_local_options_invalid(self, name, value): grid = Grid(shape=(8, 8, 8)) u = TimeFunction(name='u', grid=grid) with pytest.raises(InvalidOperator, match=f'Illegal `{name}` value'): Operator(Eq(u.forward, u + 1), opt=opt0({name: value})) def test_custom_block(self): grid = Grid(shape=(8, 8, 8)) u = TimeFunction(name='u', grid=grid) tb = (2, 2, 2) rule = None op = Operator(Eq(u.forward, u.dx + 1), opt=opt0({'par-tile': (tb, rule, 'kernel0')})) # Check generated code launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == 1 assert launches[0].block.cargs == tb assert launches[0].shm == 0 def test_large_blocks(self): """ Make sure the CUDA thread block limits are honoured (e.g., max number of threads along Z dimension = 64). """ grid = Grid(shape=(4, 4, 4)) u = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u + 1) make_opt = lambda i: opt0({'par-tile': (i, None, 'kernel0')}) # OK cases: x <= 1024, y <= 1024, x*y*z <= 1024, z <= 64 for i in [(1024, 1, 1), (1, 1024, 1), (8, 4, 32), (1, 4, 64), (128, 1, 1), (1, 1, 1)]: u.data[:] = 0. op = Operator(eq, opt=make_opt(i)) op.apply(time_M=9) assert np.all(u.data[1] == 9.) # CUDA-breaking cases -- the negations of the ones above, plus a # few others for i in [(1025, 1, 1), (1, 1025, 1), (8, 8, 32), (1, 4, 128), (1, 1, 1, 1)]: with pytest.raises(InvalidOperator): op = Operator(eq, opt=make_opt(i)) def test_device_symbols(self): grid = Grid(shape=(4, 4)) u = TimeFunction(name='u', grid=grid) op = Operator(Eq(u.forward, u + 1), opt=opt0()) assert len([i for i in op.parameters if isinstance(i, DeviceRM)]) == 1 assert len([i for i in op.parameters if isinstance(i, DeviceID)]) == 1 def test_efficient_C_land_jump(self): grid = Grid(shape=(8, 8)) u = TimeFunction(name='u', grid=grid) v = TimeFunction(name='v', grid=grid) w = TimeFunction(name='w', grid=grid) eqns = [Eq(u.forward, u + 1), Eq(v.forward, v + 1), Eq(w.forward, w + 1)] op = Operator(eqns, opt=opt0()) op.cfunction assert len(op.body.maps) == 7 assert 'Synchronize' in str(op.body.maps[3]) def test_avoid_void_subdomain_launch(self): """ CUDA fails when attempting to launch a kernel over an empty SubDomain. This test ensures we avoid the launch in such cases. """ class sd0(SubDomain): name = 'subdomain' def define(self, dimensions): x, y = dimensions return {x: ('middle', 4, 4), y: ('middle', 4, 4)} s_d0 = sd0() grid = Grid(shape=(8, 8), subdomains=(s_d0,)) u = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u + 1, subdomain=s_d0) op = Operator(eq, opt=opt0({'errctl': 'basic'})) op.apply(time_M=9) assert np.all(u.data == 0.) @pytest.mark.parallel(mode=[(4, 'diag2')]) def test_avoid_void_sparse_launch(self, mode): """ CUDA fails when attempting to launch a kernel over an empty sparse set. This test ensures we avoid the launch in such cases. """ grid = Grid(shape=(8, 8), extent=(7.0, 7.0)) x, y = grid.dimensions f = Function(name='f', grid=grid) s = SparseFunction(name='s', grid=grid, npoint=1, coordinates=[(0.5, 0.5)]) eq = s.inject(field=f, expr=s + 1) op = Operator(eq, opt=opt0({'errctl': 'basic'})) # As long as no error is raised, we're good op.apply() @pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) def test_complex(self, dtype): grid = Grid((5, 5)) x, y = grid.dimensions u = Function(name="u", grid=grid, dtype=dtype) eq = Eq(u, x + sympy.I*y + exp(sympy.I + x.spacing)) op = Operator(eq) op() # Check against numpy dx = grid.spacing_map[x.spacing] xx, yy = np.meshgrid(np.linspace(0, 4, 5), np.linspace(0, 4, 5)) npres = xx + 1j*yy + np.exp(1j + dx) assert np.allclose(u.data, npres.T, rtol=1e-7, atol=0) def setup_basic(self, dtype): grid = Grid(shape=(5,), extent=(4.,)) f = Function(name='f', grid=grid, dtype=dtype) f.data_with_halo[:] = np.arange(7) + 1j*np.arange(7, 14)[::-1] f_real = Function(name='f_real', grid=grid) f_imag = Function(name='f_imag', grid=grid) return f, f_real, f_imag def test_complex_parts_trivial(self): f, f_real, f_imag = self.setup_basic(np.complex64) eq_re = Eq(f_real, Real(f+1.)) eq_im = Eq(f_imag, Imag(f+1.)) Operator([eq_re, eq_im])() rcheck = np.array([2., 3., 4., 5., 6.]) icheck = np.array([12., 11., 10., 9., 8.]) assert np.all(np.isclose(f_real.data, rcheck)) assert np.all(np.isclose(f_imag.data, icheck)) def test_complex_parts_imag(self): f, f_real, f_imag = self.setup_basic(np.complex64) eq_re = Eq(f_real, Real(f+1j)) eq_im = Eq(f_imag, Imag(f+1j)) Operator([eq_re, eq_im])() rcheck = np.array([1., 2., 3., 4., 5.]) icheck = np.array([13., 12., 11., 10., 9.]) assert np.all(np.isclose(f_real.data, rcheck)) assert np.all(np.isclose(f_imag.data, icheck)) # Currently failing with HIP; hip_complex types required, but no # arithmetic, type conversions etc defined for these types. @skipif('nocuda') @pytest.mark.parametrize('dim', ['x', 'y']) def test_complex_parts_deriv(self, dim): # Note: needs to be 2D, as beta=1 will try (and fail) # to turn the x loop into a kernel in 1D grid = Grid(shape=(5, 5), extent=(4., 4.)) f = Function(name='f', grid=grid, dtype=np.complex64) if dim == 'x': f.data_with_halo[:] = (np.arange(7)[:, np.newaxis] + 1j*np.arange(7, 14)[::-1, np.newaxis]) else: f.data_with_halo[:] = (np.arange(7)[np.newaxis, :] + 1j*np.arange(7, 14)[np.newaxis, ::-1]) f_real = Function(name='f_real', grid=grid) f_imag = Function(name='f_imag', grid=grid) if dim == 'x': eq_re = Eq(f_real, Real(f.dx)) eq_im = Eq(f_imag, Imag(f.dx)) else: eq_re = Eq(f_real, Real(f.dy)) eq_im = Eq(f_imag, Imag(f.dy)) Operator([eq_re, eq_im])() assert np.all(np.isclose(f_real.data, 1.)) assert np.all(np.isclose(f_imag.data, -1.)) def test_complex_parts_conj(self): grid = Grid(shape=(5,)) f = Function(name='f', grid=grid, dtype=np.complex64) g = Function(name='g', grid=grid, dtype=np.complex64) f.data[:] = np.arange(5) + 1j*np.arange(5)[::-1] Operator([Eq(g, Conj(f))])() assert np.all(np.isclose(g.data, np.conj(f.data))) @skipif('nocuda') # Complex atomic for HIP not yet implemented def test_complex_source(self): grid = Grid(shape=(5,), extent=(4.,)) f = Function(name='f', grid=grid, dtype=np.complex64) g = Function(name='g', grid=grid, dtype=np.complex64) sr = SparseFunction(name='sr', grid=grid, npoint=1) sc = SparseFunction(name='sc', grid=grid, npoint=1, dtype=np.complex64) sr.coordinates.data[:] = 2. sc.coordinates.data[:] = 2. src0 = sc.inject(f, 1+1j) src1 = sr.inject(g, 1) op = Operator([src0, src1]) op() assert np.all(np.isclose(f.data, np.array([0., 0., 1.+1.j, 0., 0.]))) assert np.all(np.isclose(g.data, np.array([0., 0., 1., 0., 0.]))) def test_mixed_prec_comparisons(self): grid = Grid(shape=(11, 11, 11)) f = Function(name='f', grid=grid, dtype=np.float32) g = Function(name='g', grid=grid, dtype=np.float16) h = Function(name='h', grid=grid, dtype=np.int32) args = [(f, g), (g, f), (f, Cast(1., dtype=np.float16))] ops = [Gt, Lt, Ge, Le, Ne, CondEq] eqs = [Eq(h, INT(op(*arg))) for arg, op in product(args, ops)] op = Operator(eqs) op() class SpeccedDomain(SubDomain): """ A SubDomain constructed from one side specification per Dimension. """ separated = True def __init__(self, name, spec, thickness, **kwargs): self.name = name self.spec = spec self.thickness = thickness super().__init__(**kwargs) def define(self, dimensions): mapper = {} for dimension, side in zip(dimensions, self.spec, strict=True): if side == 'middle': mapper[dimension] = (side, self.thickness, self.thickness) elif side == 'none': mapper[dimension] = dimension else: mapper[dimension] = (side, self.thickness) return mapper class TestAliases: def test_sin_cos_v0(self): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid) g = Function(name='g', grid=grid) u = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u + sin(f) + cos(g)) op = Operator(eq, opt=opt0({'cire-mingain': 100})) kernel = op._func_table['kernel0'].root expr = kernel.body.body[-1].body[0].body[0].expr cexpr = op._printer().doprint(expr, None) assert all(i in cexpr for i in ('__cosf', '__sinf')) def test_sin_cos_v1(self, host_kwargs): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid) g = Function(name='g', grid=grid) u = TimeFunction(name='u', grid=grid) u1 = TimeFunction(name='u', grid=grid) f.data[:] = .01 g.data[:] = .03 eq = Eq(u.forward, u + sin(f) + cos(f) + sin(g) + cos(g)) op0 = Operator(eq, **host_kwargs) op1 = Operator(eq, opt=opt0({'cire-mingain': 100})) assert '__sincos' in str(op1) assert not any(i in str(op1) for i in ('sin(', '__sin(', ' cos(', ' __cos(')) op0.apply(time_M=5) op1.apply(time_M=5, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def test_sincos_host_not_intrinsic(self): """ Sin/cos hoisted to the host prelude must not use the device-only `__sinf`/`__cosf`/`__sincosf` intrinsics. """ grid = Grid(shape=(16, 16)) time = grid.time_dim u = TimeFunction(name='u', grid=grid, time_order=2, space_order=4) # Argument depends only on `time`, so it is computed on the host wt = 2 * np.pi * 0.01 * time * time.spacing eq = Eq(u.forward, 2 * u - u.backward + u.laplace + cos(wt) - sin(wt)) op = Operator([eq], opt=opt0()) cstr = str(op) assert not any(i in cstr for i in ('__sinf', '__cosf', '__sincos')) assert 'cos(' in cstr and 'sin(' in cstr def test_time_invariant(self, host_kwargs): grid = Grid(shape=(10, 10)) f = TimeFunction(name='f', grid=grid) g = Function(name='g', grid=grid) eq = Eq(f.forward, f + cos(g)) op = Operator(eq, opt=opt0({'cire-mingain': 50})) # Check generated code # NOTE: Expected one launch outside the time loop (due to lifting), # and one inside assert len(op._func_table) == 2 + 4 # +4 -> alloc_on_device/map/unmaps section0 = op.body.body[-1].body[0] timeloop = op.body.body[-1].body[1] assert timeloop.is_Iteration and timeloop.dim is grid.time_dim assert len(FindNodes(KernelLaunch).visit(section0)) == 1 assert len(FindNodes(KernelLaunch).visit(timeloop)) == 1 op.apply(time_M=9) assert np.all(f.data[0] == 10.) @pytest.mark.skip("Still unable to properly schedule space invariants") def test_space_invariant(self, host_kwargs): grid = Grid(shape=(10, 10, 10)) x, y, z = grid.dimensions f = Function(name='f', grid=grid) eq = Eq(f, f + cos(x*y*z) + cos(x*y)*cos(y) + sin(x)*cos(x*z)) op0 = Operator(eq, **host_kwargs) # noqa op1 = Operator(eq, opt=opt0()) # noqa def test_sincos_deriv_unhoisted(self): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid) u1 = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u + sin(f).dx + cos(f).dx) op0 = Operator(eq, opt=opt0()) op1 = Operator(eq, opt=opt0({'cire-mingain': np.inf})) init_chessboard(f) op0(time_M=5) op1(time_M=5, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def test_sin_cos_hoisted_nontrivial(self): """ Test that temporaries for sin and cos are created, but no compound expressions involving both sin and cos are created, for example, `sin(f)` and `cos(f)` are hoisted but not `sin(f)*cos(f)`. """ grid = Grid(shape=(40, 30, 20)) space_order = 8 u = TimeFunction(name='u', grid=grid, space_order=space_order, time_order=1, save=Buffer(1)) f = Function(name='f', grid=grid, space_order=space_order) g = f.func(name='g') h = f.func(name='h') R = [[1-(1-cos(f))*cos(g)**2, -(1-cos(f))*sin(g)*cos(g), -sin(f)*cos(g)], [-(1-cos(f))*sin(g)*cos(g), 1-(1-cos(f))*sin(g)**2, -sin(f)*sin(g)], [sin(f)*cos(g), sin(f)*sin(g), cos(f)]] R = np.array(R) expr = (h*(R[0, 0].dx) + (h*R[0, 1]).dx.dy + (h*R[0, 2]).dx.dz + (h*R[1, 0]).dx.dy + (h*R[1, 1]).dy + (h*R[1, 2]).dy.dz + h*R[2, 0] + h*R[2, 1] + h*R[2, 2]) eqn = Eq(u.forward, expr.dx + expr.dy + u) op = Operator(eqn) op.cfunction symbols = FindSymbols().visit(op.body) arrays = [c for i in symbols for c in (i.components if i.is_Bundle else (i,)) if isinstance(c, TempArray)] assert len(arrays) == 7 def test_2pass_guards_unoptimized_kernels(self): grid = Grid(shape=(31, 17, 25)) x, _, _ = grid.dimensions so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so) eqn = Eq(u.forward, u.dy.dx + u.dy.dz + f*u.dx.dz) op = Operator(eqn, opt=opt0({'cire-mode': '2pass'})) op.cfunction kernel0 = op._func_table['kernel0'].root assert 'kernel1' in op._func_table _, c1 = FindNodes(Conditional).visit(kernel0) assert c1.condition == sympy.And(x >= x.symbolic_min, x <= x.symbolic_max) class TestStreaming: def test_save_two_buffers(self): nt = 10 grid = Grid(shape=(4, 4)) u = TimeFunction(name='u', grid=grid) usave = TimeFunction(name='usave', grid=grid, save=nt, layers=NoLayers) vsave = TimeFunction(name='vsave', grid=grid, save=nt, layers=NoLayers) eqns = [Eq(u.forward, u + 1.), Eq(usave, u), Eq(vsave, u)] op = Operator(eqns) # Check generated code assert len(op._func_table) == 11 def test_save_w_layers(self): nt = 10 grid = Grid(shape=(4, 4)) u = TimeFunction(name='u', grid=grid) usave = TimeFunction(name='usave', grid=grid, save=nt) vsave = TimeFunction(name='vsave', grid=grid, save=nt) eqns = [Eq(u.forward, u + 1.), Eq(usave, u), Eq(vsave, u)] op = Operator(eqns) op.cfunction ws = op._func_table['write_snapshot0'].root symbols = FindSymbols().visit(ws) assert len([i for i in symbols if isinstance(i, Stream)]) == 1 assert len(op._func_table) == 17 assert 'write_snapshot0' in op._func_table assert 'write_snapshot1' not in op._func_table def test_streaming_two_buffers(self): nt = 10 grid = Grid(shape=(4, 4)) u = TimeFunction(name='u', grid=grid) usave = TimeFunction(name='usave', grid=grid, save=nt, layers=NoLayers) vsave = TimeFunction(name='vsave', grid=grid, save=nt, layers=NoLayers) eqn = Eq(u.forward, u + usave + vsave) op = Operator(eqn) # Check generated code assert len(op._func_table) == 11 assert len([i for i in FindSymbols().visit(op) if i.is_Array]) == 12 class TestAdvanced: @skipif('nocuda') @pytest.mark.parametrize('shape,offgrid,bounded', [ ((17, 19, 21), True, False), ((17, 19, 21), True, True), ((256, 256, 256), False, False), ], ids=['offgrid', 'bounded', 'plane']) @pytest.mark.parametrize('gpu_max_ilp', [False, 1, 2]) @pytest.mark.parametrize('gpu_unroll_ths', [0, 16]) def test_sinc_receivers(self, shape, offgrid, bounded, gpu_max_ilp, gpu_unroll_ths, host_kwargs): """Scalar sinc interpolation at off-grid points and on a receiver plane.""" nt = 5 grid = Grid(shape=shape, extent=tuple(20.*(n - 1) for n in shape)) p = TimeFunction(name='p', grid=grid, space_order=8, time_order=2) q = p.func(name='q') rcv = SparseTimeFunction(name='rcv', grid=grid, nt=nt, npoint=shape[0]*shape[1], interpolation='sinc', r=4) eqns = rcv.interpolate(expr=p.forward/3.0 + 2.0*q.forward/3.0) op0 = Operator(eqns, name='SincReceiversHost', **host_kwargs) # The per-nest tile form overrides the small-kernel launch heuristic op1 = Operator(eqns, name='SincReceivers', opt=('advanced', {'par-tile': ((128,), (128,)), 'gpu-reduce-local': False, 'gpu-max-ilp': gpu_max_ilp, 'gpu-unroll-ths': gpu_unroll_ths})) assert '#pragma unroll(4)' in str(op1) # Distinct oscillations across space and time, including the halo t, x, y, z = np.ogrid[tuple(slice(n) for n in p.data_with_halo.shape)] p.data_with_halo[:] = ((t + x + 2*y + 3*z) % 7 - 3)/7. q.data_with_halo[:] = ((2*t + 3*x + y + z) % 5 - 2)/5. # Populate coordinates after construction: the kernel must accept # runtime positions, including stencils that extend into the halo if offgrid: for axis, (n, h) in enumerate(zip(shape, grid.spacing)): positions = np.linspace(.25, n - 1.25, n) rcv.coordinates.data[:, axis] = np.resize(positions, rcv.npoint)*h rcv.coordinates.data[0, :] = 0 rcv.coordinates.data[1, :] = grid.extent else: x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij') rcv.coordinates.data[:, 0] = x.ravel()*grid.spacing[0] rcv.coordinates.data[:, 1] = y.ravel()*grid.spacing[1] rcv.coordinates.data[:, 2] = grid.extent[2]/2 bounds = dict(x_m=6, x_M=10, y_m=6, y_M=12, z_m=6, z_M=12) if bounded else {} op0.apply(time_m=0, time_M=nt-1, **bounds) expected = rcv.data.copy() rcv.data[:] = np.nan op1.apply(time_m=0, time_M=nt-1, **bounds) assert np.isfinite(expected).all() assert np.any(expected != 0) if bounded: assert np.all(expected[:, 1] == 0) assert np.allclose(rcv.data, expected, rtol=1e-5, atol=1e-6) def test_code_size(self): grid = Grid(shape=(4, 4)) f = Function(name='f', grid=grid) u = TimeFunction(name='u', grid=grid) op = Operator(Eq(u.forward, u + f + 1)) op.cfunction assert len(op._func_table) == 4 assert 'map_to_device0' in op._func_table assert 'map_to_device1' not in op._func_table assert 'unmap_from_device0' in op._func_table # Performs copy-back assert 'unmap_from_device1' in op._func_table # no copy-back (readonly) @pytest.mark.parallel(mode=[(1, 'diag2')]) def test_mpi_two_eqns(self, mode): grid = Grid(shape=(16, 16)) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dx + v + 1.), Eq(v.forward, v.dx + u + 1.)] op = Operator(eqns) expected = {'gather0', 'scatter0', 'haloupdate0', 'halowait0', 'kernel2', 'kernel3', 'kernel4', 'map_to_device0', 'unmap_from_device0', 'init_global0'} if op._options['gpu-vect']: expected.update(['alloc2', 'bundle2', 'unbundle1', 'free2']) assert set(op._func_table) == expected @pytest.mark.parallel(mode=[(4, 'diag2')]) def test_mpi_issue_126(self, mode): shape = (10, 10, 10) spacing = (10.0, 10.0, 10.0) origin = tuple([0.0 for s in shape]) extent = tuple([d * (s - 1) for s, d in zip(shape, spacing)]) grid = Grid(extent=extent, shape=shape, origin=origin) time_range = TimeAxis(start=0., stop=400., step=1.) p0 = TimeFunction(name='p0', grid=grid, time_order=2, space_order=4, save=Buffer(2)) src_coords = np.empty((1, len(grid.shape)), dtype=grid.dtype) src_coords[0, :] = [d * (s-1)//2 for d, s in zip(grid.spacing, grid.shape)] src_coords[:, -1] = 0. src_coords[:] = [45, 0, 0] src = RickerSource(name='src', grid=grid, f0=0.010, npoint=1, time_range=time_range) src.coordinates.data[:] = src_coords[:] eqns = [Eq(p0.forward, p0.dz.dx + p0.dz.dz + p0.dx + p0.dz), src.inject(field=p0.forward, expr=src)] op = Operator(eqns) op(time=20, dt=time_range.step) assert np.isclose(norm(p0), 0.01951365, rtol=1e-04) def test_issue_329(self): grid = Grid(shape=(10, 10)) time = grid.time_dim nf = 1 fq = DefaultDimension(name="fq", default_value=nf) f = Function(name="f", grid=grid, dimensions=(fq,), shape=(nf,)) p0 = TimeFunction(name="p0", grid=grid) p0re = Function(name="p0re", grid=grid, dimensions=(fq, *grid.dimensions), shape=(nf, *grid.shape)) eqns = [Eq(p0re, p0re + p0*cos(time.spacing*(time - 1)*f)), Eq(p0.forward, p0 + 1)] op = Operator(eqns) # Ensure it jit-compiles and runs smoothly op.apply(time_M=4, dt=.1) def test_custom_partile_w_trivial_kernel(self): grid = Grid(shape=(30, 30, 30)) x, y, z = grid.dimensions f = Function(name='f', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4) eq = Eq(u[0, x, y, z], f + 1) op = Operator(eq, opt=opt0({'par-tile': (10, 10)})) # Check generated code # Given a 3D physical domain, a 2D custom par-tile should default to # a 3D par-tile in the case of trivial kernels without a stencil launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == 1 assert launches[0].block.cargs == (32, 4, 1) assert launches[0].shm == 0 @pytest.mark.parametrize('par_tile,block', [ ((32, 4), (32, 4, 1)), (((32, 4),), (32, 4, 1)), (((32, 4), (32, 4)), (32, 4, 4)), ((16, 4), (32, 4, 1)), (((16, 4), (8, 8)), (16, 4, 4)), ((32, 4, 4), (32, 4, 1)), (((16, 4, 4), (8, 8, 8)), (16, 4, 4)), ]) def test_copy_partile_blocking(self, par_tile, block): grid = Grid(shape=(40, 41, 47)) u = Function(name='u', grid=grid, space_order=0) v = Function(name='v', grid=grid, space_order=0) op = Operator(Eq(u, v), name='copy_partile', opt=opt0({'par-tile': par_tile})) launch, = FindNodes(KernelLaunch).visit(op) assert launch.block.cargs == block assert len(launch.grid.shape) == 3 assert not any(i.dim.is_Space for i in FindNodes(Iteration).visit(op)) expected = np.arange(np.prod(grid.shape), dtype=np.float32).reshape(grid.shape) v.data[:] = expected u.data[:] = -1 op.apply() np.testing.assert_array_equal(u.data, expected) @pytest.mark.parametrize('par_tile,xchange,blocks,grid_ranks', [ (((16, 4), (8, 8)), True, [(16, 4, 4), (8, 8)], [3, 2]), (((16, 4, 4), (8, 8, 4)), False, [(16, 4, 4), (8, 8, 4)], [3, 3]), ]) def test_copy_stencil_partile_mapping(self, par_tile, xchange, blocks, grid_ranks): grid = Grid(shape=(40, 41, 47)) u = Function(name='u', grid=grid, space_order=2) v = Function(name='v', grid=grid, space_order=2) w = Function(name='w', grid=grid, space_order=2) u1 = u.func(name='u1') w1 = w.func(name='w1') eqns = [Eq(u, v), Eq(w, u.dx)] op = Operator(eqns, name='copy_stencil_partile', opt=opt0({'par-tile': par_tile}, xchange=xchange)) reference = Operator(eqns, name='copy_stencil_reference', opt=opt0({'par-tile': (32, 4, 4)})) launches = FindNodes(KernelLaunch).visit(op) assert [i.block.cargs for i in launches] == blocks assert [len(i.grid.shape) for i in launches] == grid_ranks assert not any(i.dim.is_Space for i in FindNodes(Iteration).visit(op)) x = grid.dimensions[0] kernel = op._func_table['kernel1'].root xloops = [i for i in FindNodes(Iteration).visit(kernel) if i.dim.root is x] if grid_ranks[1] == 2: xloop, = xloops assert xloop.limits == (x.symbolic_min, x.symbolic_max, 1) else: assert not xloops v.data[:] = np.arange(np.prod(grid.shape), dtype=np.float32).reshape(grid.shape) reference.apply(u=u1, w=w1) op.apply() np.testing.assert_array_equal(u.data, v.data) np.testing.assert_allclose(w.data, w1.data) def test_custom_partile_w_custom_dim(self): grid = Grid(shape=(10, 10, 10)) x, y, z = grid.dimensions zc = CustomDimension(name='zc', parent=z, symbolic_min=1, symbolic_max=8) u = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u + 1).subs(z, zc) op = Operator(eq, opt=opt0({'par-tile': (16, 16)})) launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == 1 # Small kernel => ignore the custom par-tile assert launches[0].block.cargs == (32, 4, 1) assert launches[0].shm == 0 def test_fissioning_x_for_parallelism(self): grid = Grid((11, 11, 11)) x, y, z = grid.dimensions image_vs = Function(name='image_vs', grid=grid, space_order=1, staggered=NODE) p_back_xy = TimeFunction(name='p_back_xy', grid=grid, staggered=(x, y), space_order=4, time_order=1, save=Buffer(1)) eqns = [Eq(image_vs, p_back_xy + image_vs), Eq(p_back_xy.backward, p_back_xy)] op = Operator(eqns) launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == 2 assert len(launches[0].grid.cargs) == 3 assert len(launches[1].grid.cargs) == 3 # Make sure there is no serial x loop assert 'for (int x = x_m; x <= x_M; x += 1)' not in str(op) class TestKernelOptInterchange: def _check_interchange(self, op, grid, expected): x, _, _ = grid.dimensions # Check time loop wraps the kernel launch trees = retrieve_iteration_tree(op) assert len(trees) == 1 tree = trees.pop() assert len(tree) == 1 assert tree[0].dim is grid.time_dim # Check kernel loop nest efunc = op._func_table['kernel0'].root trees = retrieve_iteration_tree(efunc) assert len(trees) == expected assert len(trees[0]) == 1 assert all(x in i[0].dim._defines for i in trees) def test_basic(self): grid = Grid(shape=(51, 33, 42)) u = TimeFunction(name='u', grid=grid, space_order=4) op = Operator(Eq(u.forward, u.dx + 1), opt=opt1({'par-tile': (32, 4)})) # Check generated code # Expected: # kernel_launch<<...>> # # __global__ kernel0(): # for x # ... self._check_interchange(op, grid, 2) op.cfunction def test_time_parallel(self): grid = Grid(shape=(31, 17, 25)) f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4, save=5, layers=NoLayers) v = u.func(name='v') eqns = [Eq(v.forward, f*u.dy.dx + g*u.dy.dz)] op = Operator(eqns, opt=opt1({'gpu-fit': (u, v), 'par-tile': (16, 4, 4)})) # Check generated code self._check_interchange(op, grid, 3) op.cfunction def test_2pass_guards_normal_form(self, host_kwargs): grid = Grid(shape=(31, 17, 25), extent=(30., 16., 24.)) so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so) u1 = u.func(name='u1') eqn = Eq(u.forward, u + .001*(u.dy.dx + u.dy.dz + f*u.dx.dz)) op0 = Operator(eqn, name='GuardedCIREHost', **host_kwargs) op1 = Operator(eqn, name='GuardedCIRE', opt=opt1({ 'cire-mode': '2pass', 'par-tile': (32, 4, 4) })) # Keep a two-dimensional thread block despite the helper's distinct guards assert FindNodes(KernelLaunch).visit(op1)[0].block.cargs == (32, 4) f.data_with_halo[:] = 1.5 init_chessboard(u) init_chessboard(u1) op0.apply(time_M=2, u=u1) op1.apply(time_M=2) assert np.allclose(u.data, u1.data, rtol=1.e-5, atol=1.e-6) class TestKernelOptRegRot: @pytest.mark.parametrize('expr', [ 'u.dx', 'u.dx2', 'u.dx + u.dx2', ]) def test_regrot(self, expr, host_kwargs): grid = Grid(shape=(51, 33, 42)) t = grid.stepping_dim # noqa x, y, z = grid.dimensions # noqa f = Function(name='f', grid=grid) # noqa u = TimeFunction(name='u', grid=grid, space_order=4, time_order=2) u1 = TimeFunction(name='u', grid=grid, space_order=4, time_order=2) eq = Eq(u.forward, eval(expr) + 1) op0 = Operator(eq, **host_kwargs) op1 = Operator(eq, opt=opt2({'par-tile': (32, 4)})) # Check generated code assert 'queue' in str(op1) efunc = op1._func_table['kernel0'].root trees = retrieve_iteration_tree(efunc) assert len(trees) == 4 assert len(trees[0]) == 1 assert isinstance(trees[0][0].dim, StencilDimension) assert len(trees[1]) == 2 assert trees[1][0].dim is x assert isinstance(trees[1][1].dim, StencilDimension) op0.apply(time_M=9) op1.apply(time_M=9, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e5) @pytest.mark.parametrize('axis', [1, 2]) def test_harmonic_interpolation(self, axis, host_kwargs): grid = Grid(shape=(17, 13, 19)) x = grid.dimensions[0] f = Function(name='f', grid=grid, space_order=4, avg_mode='safe_harmonic') g = Function(name='g', grid=grid, space_order=4, staggered=(x, grid.dimensions[axis])) g1 = g.func(name='g1') eq = Eq(g, f) op0 = Operator(eq, name='op0', **host_kwargs) op1 = Operator(eq, name='op1', opt=opt4({'par-tile': (32, 4)})) # The one-element late queue is fully primed at physical index zero assert 'float queue1[1];' in str(op1) assert 'queue1[qq0] = fL0' in str(op1) assert 'SAFEINV(queue1[qq0], queue1[qq0])' in str(op1) f.data_with_halo[:] = 3.5 f.data[:] = np.arange(grid.shape[0])[:, None, None] + 3.5 op0.apply() op1.apply(g=g1) assert np.allclose(g.data, g1.data, rtol=1.e-6) def test_regrot_queue_dtype_with_mixed_dtypes(self): """ The register queue dtype must match the dtype of the expression it stores, not the enclosing kernel dtype. Regression test for a bug where a real TimeFunction under register rotation got a complex-typed queue because a sibling complex Function promoted `kernel.dtype` to complex, producing uncompilable code such as `s_y0[...] = queue0[...]` (float = thrust::complex). """ grid = Grid(shape=(64, 64)) time = grid.time_dim nfreq = 4 v = TimeFunction(name='v', grid=grid, time_order=2, space_order=8) uf = Function(name='uf', dimensions=(Dimension('freq_dim'),) + grid.dimensions, grid=grid, shape=(nfreq,) + grid.shape, dtype=np.complex64) dt = time.spacing src = sum(uf._subs(uf.indices[0], i) * exp(sympy.I * 2 * np.pi * 0.01 * (i + 1) * time * dt) for i in range(nfreq)) eq = Eq(v.forward, 2 * v - v.backward + v.laplace + Real(src)) op = Operator(eq, opt=opt4({'par-tile': (32, 4)})) kernel = op._func_table['kernel0'].root queues = [i for i in FindSymbols().visit(kernel) if isinstance(i, TBQueue)] assert queues, "expected at least one register queue" for q in queues: assert q.dtype == v.dtype, \ f"queue {q.name} dtype {q.dtype} != v.dtype {v.dtype}" # And the operator must actually build (the original bug crashed nvcc) op.cfunction class TestKernelOptShm: @pytest.mark.parametrize('shape', [(5, 7, 21), (40, 39, 38), (21, 11, 89), (20, 10, 40)]) @pytest.mark.parametrize('par_tile', [(4, 4, 17), (6, 5, 4), (32, 4, 4)]) def test_simple(self, shape, par_tile, host_kwargs): grid = Grid(shape=shape) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) init_chessboard(u) init_chessboard(u1) eqn = Eq(u.forward, u.dx.dy + 1.) op0 = Operator(eqn, **host_kwargs) op1 = Operator(eqn, opt=opt3({'par-tile': par_tile, 'gpu-loader-threads': '1D'})) # Check generated code assert 's_yz' in str(op1) op0.apply(time_M=5) op1.apply(time_M=5, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) @pytest.mark.parametrize('shape', [(5, 7, 21), (40, 39, 38), (21, 11, 89)]) def test_complex(self, shape, host_kwargs): grid = Grid(shape=shape) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) u2 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) v1 = TimeFunction(name='v', grid=grid, space_order=4) v2 = TimeFunction(name='v', grid=grid, space_order=4) eqns = [ Eq(u.forward, u.dx.dy + u.dy.dz + u.dx.dz + v + 1.), Eq(v.forward, v.dx.dy + v.dy.dz + v.dx.dz + u + 1.) ] op0 = Operator(eqns, **host_kwargs) op1 = Operator(eqns, opt=opt0()) op2 = Operator(eqns, opt=opt3({'par-tile': (32, 4, 4)})) # Check generated code -- the peculiar structure of the device kernel assert 'queue' not in str(op2) assert 's_yz0' in str(op2) and 's_yz0' not in str(op1) assert 's_yz1' in str(op2) and 's_yz1' not in str(op1) op0.apply(time_M=1) op1.apply(time_M=1, u=u1, v=v1) op2.apply(time_M=1, u=u2, v=v2) assert np.allclose(u.data, u1.data, rtol=1.e-4) assert np.allclose(v.data, v1.data, rtol=1.e-4) assert np.allclose(u1.data, u2.data, rtol=1.e-4) assert np.allclose(v1.data, v2.data, rtol=1.e-4) @skipif('nocuda-cc80') @pytest.mark.parametrize('offset', [0, 7, 97, 125]) def test_reuse_order(self, offset): grid = Grid(shape=(24, 24, 24)) cp, sp, st, ct, vz, wy = [ Function(name=name, grid=grid, space_order=8) for name in ('cp', 'sp', 'st', 'ct', 'vz', 'wy') ] u = TimeFunction(name='u', grid=grid, space_order=8) expr = ((st*cp*vz).dx + (st*sp*vz).dy + (ct*vz).dz + (-sp*wy).dx + (cp*wy).dy) eqn = Eq(u.forward, expr + 1) # Simulate shared arrays allocated by earlier kernels sregistry = SymbolRegistry() for _ in range(offset): sregistry.make_name(prefix='s_yz') options = {'cire-mode': '1pass', 'gpu-opt': True, 'gpu-glb-mem': 'async-pipe', 'gpu-vect': 'stream-only', 'gpu-hoist': 2, 'gpu-prefetch': 3, 'gpu-wdr': False, 'par-tile': (32, 4, 48), 'gpu-graph': False} op = Operator(eqn, name='op', opt=('advanced', options), sregistry=sregistry) kernel = op._func_table['kernel0'].root assert sum(i.name == '__syncthreads' for i in FindNodes(Call).visit(kernel)) == 8 for f in (cp, sp, st, ct, vz, wy): init_chessboard(f) reference = Operator(eqn, name='reference', opt=opt0()) reference.apply(time_M=4) expected = u.data.copy() u.data_with_halo[:] = 0 op.apply(time_M=4) assert np.allclose(u.data, expected, rtol=1.e-6, atol=1.e-6) class TestKernelOptDefault: def test_1step(self, host_kwargs): grid = Grid(shape=(40, 39, 38)) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) init_chessboard(u) init_chessboard(u1) eqn = Eq(u.forward, u.dx.dy + u.dy.dz + 1.) op0 = Operator(eqn, **host_kwargs) op1 = Operator(eqn, opt=opt4({'par-tile': (4, 4, 17), 'gpu-loader-threads': '1D', 'gpu-opt-steps': 1})) # Check generated code assert 'queue' in str(op1) op0.apply(time_M=5) op1.apply(time_M=5, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def test_to2(self): grid = Grid(shape=(32, 41, 53)) u = TimeFunction(name='u', grid=grid, space_order=4, time_order=2) u1 = TimeFunction(name='u', grid=grid, space_order=4, time_order=2) init_chessboard(u) init_chessboard(u1) eqn = Eq(u.forward, u.dx.dy + u.dy.dx + u.backward + 1.) op0 = Operator(eqn, opt=opt0()) op1 = Operator(eqn, opt=opt4({'par-tile': (16, 8)})) op0.apply(time_M=4) op1.apply(time_M=4, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-6) @pytest.mark.parametrize('npoint', [1, 17]) def test_sparse_ops(self, npoint): grid = Grid(shape=(4, 4)) f = Function(name='f', grid=grid) coords = [(1., 1.)]*npoint sf0 = SparseFunction(name='sf0', grid=grid, npoint=npoint, coordinates=coords) sf1 = SparseFunction(name='sf1', grid=grid, npoint=npoint, coordinates=coords) eqns = sf0.inject(field=f, expr=sf0 + 1) eqns += sf1.interpolate(expr=f + 1) op = Operator(eqns, opt=opt4({'par-tile': (1, 1)})) op.apply() assert np.allclose(sf1.data, npoint + 1, rtol=1.e-6) def test_safe_inv(self): grid = Grid(shape=(31, 17, 40)) x, _, _ = grid.dimensions mu = Function(name='mu', grid=grid, avg_mode='safe_harmonic') f = Function(name='f', grid=grid, staggered=x) u = TimeFunction(name='u', grid=grid, space_order=4) eqn = Eq(f, (u*mu).dy) op = Operator(eqn, opt=opt4({'par-tile': (10, 7), 'cire-mingain': 10000})) op.cfunction @pytest.mark.parametrize('expr', [ 'u.dy', 'u.dy.dx', 'u.dx.dy', 'u.dy.dy', ]) def test_2D(self, expr): grid = Grid(shape=(31, 17)) nt = 5 # We use save otherwise the machine error grows too quickly u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) eqn = Eq(u.forward, eval(expr) + 1.) op0 = Operator(eqn, opt=opt0({'gpu-fit': [u]})) op1 = Operator(eqn, opt=opt4({'par-tile': (10,), 'gpu-fit': [u]})) # Check generated code if ".dx" in expr: assert 'queue' in str(op1) assert 's_y' in str(op1) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data[1], u1.data[1], rtol=1.e-5) assert np.allclose(u.data[2:], u1.data[2:], rtol=1.e-2) def test_2D_laplacian_via_queue(self): grid = Grid(shape=(11, 11)) u = TimeFunction(name='u', grid=grid, space_order=4) eqn = Eq(u.forward, u.dx2 + u.dy2) op = Operator(eqn, opt=opt4({'par-tile': (10, 5)})) # Check generated code assert 's_y0[ty + 2] = uL0(t0, x + 4, y + 4)' not in str(op) assert 's_y0[ty + 2] = queue0[2]' in str(op) @pytest.mark.parametrize('expr', [ 'u.laplace', 'u.laplace + u', ]) @pytest.mark.parametrize('par_tile', [(10, 7), (10, 7, 3)]) def test_3D_basic(self, expr, par_tile): grid = Grid(shape=(31, 17, 25)) nt = 5 u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) eqn = Eq(u.forward, eval(expr) + 1.) op0 = Operator(eqn, opt=opt0({'gpu-fit': [u]})) op1 = Operator(eqn, opt=opt4({'par-tile': par_tile, 'gpu-fit': [u]})) # Check generated code -- expected only ONE shared memory array assert len([i for i in FindSymbols().visit(op1) if i._mem_shared]) == 1 assert 'block0(10,7)' in str(op1) assert 's_yz0[ty + 2][tz + 2] = queue0[2]' in str(op1) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def _check_kernel_opt(self, op1, f, exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps=None, kids=None, ctx=None): exp_shmas = as_tuple(exp_shmas) exp_queues = as_tuple(exp_queues) exp_syncs = as_tuple(exp_syncs) exp_ngas = as_tuple(exp_ngas) assert len(exp_shmas) == len(exp_queues) == len(exp_syncs) == len(exp_ngas) exp_tmps = as_tuple(exp_tmps) or (0,)*len(exp_shmas) kids = as_tuple(kids) or tuple(range(len(exp_shmas))) test = False for n, (esh, eq, esy, en, et) in enumerate(zip(exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps)): try: kid = kids[n] except IndexError: kid = n name = 'kernel%d' % kid efunc = op1._func_table[name].root shmas = [i for i in FindSymbols().visit(efunc) if i._mem_shared] assert len(shmas) == esh queues = [i for i in FindSymbols().visit(efunc) if i.name.startswith('queue')] assert len(queues) == eq syncs = [i for i in FindNodes(Call).visit(efunc) if i.name == '__syncthreads'] assert len(syncs) == esy indexeds = FindSymbols('indexeds').visit(efunc) ngas = {i for i in indexeds if i.function is f} assert len(ngas) == en if exp_tmps: writes = [i.write for i in FindNodes(Expression).visit(efunc)] tmps = [i for i in writes if isinstance(i, TempArray)] assert len(tmps) == et assert all(i in str(efunc) for i in as_tuple(ctx)) test = True assert test @pytest.mark.parametrize('exprs,exp_shmas,exp_queues,exp_syncs,exp_ngas,reltol', [ # Single eqn simple ('u.laplace', 1, 1, 2, 7, False), ('u.laplace + u', 1, 1, 2, 7, False), # Single eqn cross-derivatives ('u.dy.dz', 2, 0, 3, 4, False), ('u.dx.dz', 1, 1, 2, 3, False), ('u.dx.dz + u.dy.dz', 2, 1, 3, 5, False), ('u.dy.dx', 1, 1, 4, 7, False), ('u.dx.dz + u', 1, 1, 2, 3, False), ('u.dx.dz + u.dy', 2, 1, 2, 5, False), ('u.dx.dz + u.dy.dx', 2, 2, 4, 7, False), ('u.dx.dz + u.dx.dy', 1, 1, 2, 3, False), ('u.dy.dz + u.dy.dx', 1, 1, 6, 7, False), ('u.dy.dz + u.dx.dz + u.dy.dx', 1, 2, 6, 7, False), ('u.dx.dy + u.dy.dz + u.dz.dx', 3, 2, 6, 3, False), ('v.dy.dz + v1.dy.dz2', 3, 0, 4, 1, False), # Single eqn cross-derivatives w/ extra fields ('(u.dx + f).dy', 1, 1, 2, 3, False), ('(f*u.dx + g*u.dy).dz', 2, 1, 3, 5, False), # Single eqn cross-derivatives + snd-order derivatives ('u.dx2 + u.dx.dy', 1, 1, 2, 3, False), ('u.dy2 + u.dx.dy', 2, 1, 2, 3, False), ('u.dz2 + u.dx.dy', 2, 1, 2, 5, True), ('u.dy2 + u.dz.dy', 2, 0, 3, 4, True), ('u.dy2 + u.dy.dx', 1, 2, 4, 7, True), ('u.dy2 + u.dy.dz', 2, 0, 3, 4, True), ('u.dx2 + u.dy2 + u.dx.dy', 2, 1, 2, 3, False), ('u.dz2 + u.dy2', 1, 0, 2, 6, False), ('u.dz2 + u.dy2 + u.dy.dz', 2, 0, 3, 4, False), ('u.dy2 + u.dx2 + u.dy.dx + u.dy.dz', 2, 3, 4, 7, True), ('u.dx + u.dy.dz + u.dy.dx + f', 1, 3, 8, 7, True), ('u.dx2 + u.dy2 + u.dz2 + u.dx.dy + u.dy.dz + f*u.dx.dz', 3, 1, 3, 3, True), ('u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz', 4, 3, 4, 7, True), ('u.dy.dz + u.dz.dy', 3, 0, 3, 2, True), ('f*(u.dy2 + v.dz2) + u.dy.dx + v.dy.dx', 2, 3, 4, 7, True), ('(u.dy2 + v.dy).dz', 3, 0, 3, 4, True), # Single eqn cross-derivatives about the same Dimension ('u.dz.dz', 1, 0, 2, 4, False), ('u.dz.dz + u.dy.dy', 1, 0, 2, 6, False), ('u.dx.dx', 0, 1, 0, 3, False), # Coupled eqns (['u.dy.dz', 'v.dy.dz'], 3, 0, 4, 4, False), (['u.dy.dz + u.dx.dz', 'v.dy.dz + v.dx.dz'], 3, 2, 4, 5, False), (['u.dy.dz + u.dx.dz + u.dy.dx', 'v.dy.dz + v.dx.dz + v.dy.dx'], 2, 4, 6, 7, False), (['u.dx.dy + u.dx.dz', 'v.dx.dy + v.dx.dz'], 2, 2, 2, 3, False), (['u.dy.dz + u.dx.dz + f*u.dy.dx', 'v.dy.dz + v.dx.dz + f*v.dy.dx'], 2, 4, 6, 7, False), (['u.dz2 + u.dy.dz', 'v.dz2 + v.dy.dz'], 4, 0, 3, 4, False), (['u.dx2 + u.dy2 + u.dz2 + u.dx.dy + u.dy.dz + f*u.dx.dz', 'v.dx2 + v.dy2 + v.dz2 + v.dx.dy + v.dy.dz + f*v.dx.dz'], 6, 2, 3, 3, True), (['u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz', 'v.dx2 + v.dy2 + v.dz2 + v.dy.dx + v.dy.dz + f*v.dx.dz'], 8, 6, 4, 7, True), ]) @pytest.mark.parametrize('subdomain', [None, 'interior']) def test_3D(self, exprs, exp_shmas, exp_queues, exp_syncs, exp_ngas, reltol, subdomain): grid = Grid(shape=(31, 17, 25)) subdomain = grid.subdomains.get(subdomain) nt = 5 f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) v = TimeFunction(name='v', grid=grid, space_order=4, save=nt, layers=NoLayers) v1 = v.func(name='v1') # noqa f.data[:] = .01 g.data[:] = .02 exprs = as_tuple(exprs) eqns = [] fields = [] for expr, field in zip(exprs, [u, v]): eqns.append(Eq(field.forward, eval(expr) + 1., subdomain=subdomain)) fields.append(field) # TODO: Horrible, but I don't wanna change the par-tile for all other # >35 tests right now... if any(e.find(u.dy.dy) or e.find(u.dz.dz) for e in eqns): par_tile = (10, 8) else: par_tile = (10, 7) op0 = Operator(eqns, opt=opt0({'gpu-fit': (u, v)})) op1 = Operator(eqns, opt=opt4({'gpu-fit': (u, v), 'par-tile': par_tile, # Legacy behavior: 'cire-hoist': False, 'cire-mode': '1pass', 'deriv-collect': True})) # Check generated code # NOTE: these are expected to be all 1-pass implementations assert 'kernel0' in op1._func_table assert 'kernel1' not in op1._func_table self._check_kernel_opt(op1, u, exp_shmas, exp_queues, exp_syncs, exp_ngas) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) if reltol: assert np.allclose(u.data, u1.data, rtol=1.e-4) else: assert np.allclose(u.data, u1.data, rtol=1.e-5) @pytest.mark.parametrize('exprs,exp_shmas,exp_queues,exp_syncs,exp_ngas,ctx', [ ('u.dy.dz + u.dz.dy', 3, 0, 3, 4, None), ('u.dz2 + u.dx.dy + u.dx.dz', 2, 1, 2, 7, None), ('u.dx2 + u.dy2 + u.dz2 + u.dx.dy + u.dy.dz + f*u.dx.dz', 2, 3, 8, 11, None), ('u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz', 4, 3, 4, 7, ['if (qq0 >= 2)', 'queue1[0]']), ('u.dy.dy + u.dz.dz', 1, 0, 2, 6, None) ]) def test_3D_lt1D(self, exprs, exp_shmas, exp_queues, exp_syncs, exp_ngas, ctx): grid = Grid(shape=(31, 17, 25)) nt = 5 f = Function(name='f', grid=grid) u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) v = TimeFunction(name='v', grid=grid, space_order=4, save=nt, layers=NoLayers) f.data[:] = .01 exprs = as_tuple(exprs) eqns = [] fields = [] for expr, field in zip(exprs, [u, v]): eqns.append(Eq(field.forward, eval(expr) + 1.)) fields.append(field) op0 = Operator(eqns, opt=opt0({'gpu-fit': fields})) op1 = Operator(eqns, opt=opt4({'par-tile': (10, 8), 'gpu-loader-threads': '1D', 'gpu-fit': fields})) # Check generated code self._check_kernel_opt(op1, u, exp_shmas, exp_queues, exp_syncs, exp_ngas, ctx=ctx) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) @pytest.mark.parametrize('exprs,exp_shmas,exp_queues,exp_syncs,exp_ngas,exp_tmps', [ ('u.dy.dz + u.dz.dy', (1, 2), (0, 0), (2, 2), (5, 1), (2, 0)), ('(u.dy + f).dz + (u.dy + g).dy', (1, 2), (0, 0), (2, 2), (3, 1), (1, 0)), (('u.dy.dx + u.dy.dz + f*u.dx.dz', 'v.dy.dx + v.dy.dz + f*v.dx.dz'), (1, 4), (2, 2), (4, 2), (4, 1), (4, 0)), ]) @pytest.mark.parametrize('shape', [(31, 17, 25), (100, 37, 50)]) def test_3D_2pass(self, exprs, exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps, shape): grid = Grid(shape=shape) nt = 5 f = Function(name='f', grid=grid, space_order=16) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=16, save=nt, layers=NoLayers) v = u.func(name='v') f.data[:] = .01 g.data[:] = .02 exprs = as_tuple(exprs) eqns = [] fields = [] for expr, field in zip(exprs, [u, v]): eqns.append(Eq(field.forward, eval(expr) + 1.)) fields.append(field) op0 = Operator(eqns, opt=opt0({'gpu-fit': fields})) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'gpu-fit': fields, 'par-tile': (10, 8), 'cire-mode': '2pass'})) # Check generated code self._check_kernel_opt( op1, u, exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps ) u1 = u.func(name='u1') init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-3) def test_1pass_and_2pass_same_op(self): grid = Grid(shape=(31, 17, 25)) u = TimeFunction(name='u', grid=grid, space_order=8) v = TimeFunction(name='v', grid=grid, space_order=16) eqns = [Eq(u.forward, u.dy.dz + v + 1), Eq(v.forward, u.forward.dx + v.dy.dz + 1)] op = Operator(eqns, opt=opt4({'par-tile': (10, 8), 'cire-mode': ('1pass', '2pass')})) # Check generated code launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == 3 key = lambda launch: [i for i in launch.writes if isinstance(i, TempArray)] assert len(key(launches[0])) == 0 assert len(key(launches[1])) == 1 assert len(key(launches[2])) == 0 def test_gen_2pass_not_3pass(self): grid = Grid(shape=(34, 45, 50)) x, y, z = grid.dimensions a = Function(name='a', grid=grid, space_order=8) vx = TimeFunction(name='vx', grid=grid, space_order=8, save=Buffer(1)) tzz = vx.func(name='tzz') x0m = x - x.spacing/2 y0p = y + y.spacing/2 z0p = z + z.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(2, 2)) expr0 = Derivative(a, y, z, x0={y: y0p, z: z0p}, **kw) expr1 = (Derivative(vx, x, y, x0={x: x0m, y: y0p}, **kw) + Derivative(vx, x, z, x0={x: x0m, z: z0p}, **kw)) eqn = Eq(tzz.forward, expr0 * expr1 + tzz) op = Operator(eqn, opt=opt4({'cire-hoist': 0, # Legacy 'cire-mode': '2pass', 'gpu-flip-derivs': True, 'par-tile': (32, 4)})) op.cfunction # Check generated code launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == 2 @pytest.mark.parametrize('exprs', [ ('u.dx + u.dy', 'v'), ('u.dx2 + u.dy + v', 'v.dx + u'), ]) @pytest.mark.parametrize('par_tile', [(2, 2), (5, 5), (10, 10)]) def test_2D_custom_partile(self, exprs, par_tile): """ Unlike several other tests, here the par-tile has as many entries as the grid shape, which will lead to a SEQUENTIAL BlockDimension whenever register rotation kicks in. """ grid = Grid(shape=(11, 11)) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) v1 = TimeFunction(name='v', grid=grid, space_order=4) eqns = [ Eq(u.forward, eval(exprs[0])), Eq(v.forward, eval(exprs[1])), ] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt4({'par-tile': par_tile})) assert 'queue' in str(op1) assert 's_y' in str(op1) op0.apply(time_M=4) op1.apply(time_M=4, u=u1, v=v1) assert np.all(u.data == u1.data) assert np.all(v.data == v1.data) @pytest.mark.parametrize('subdomain', [None, 'interior']) def test_3D_custom_partile(self, subdomain): grid = Grid(shape=(100, 64, 25)) subdomain = grid.subdomains.get(subdomain) nt = 5 u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) eqns = [Eq(u.forward, u.dx2 + u.dy.dx + 1., subdomain=subdomain)] op0 = Operator(eqns, opt=opt0({'gpu-fit': [u]})) op1 = Operator(eqns, opt=opt4({'par-tile': (64, 4, 4), 'gpu-fit': [u]})) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def test_simple_word_move_w_left_subdomain(self): class SD(SubDomain): def define(self, dimensions): x, y, z = dimensions return {x: ('left', 5), y: y, z: z} grid = Grid(shape=(10, 3, 3)) sd = SD(grid=grid) f = Function(name='f', grid=grid) u = TimeFunction(name="u", grid=grid, time_order=0) eq = Eq(u.forward, f, subdomain=sd) op = Operator(eq, opt=opt4({'errctl': 'basic', 'par-tile': (32, 4, 1)})) f.data[:] = 6 op.apply(time_M=5) assert np.all(u.data[0, :5] == 6) assert np.all(u.data[0, 5:] == 0) def test_pow_issue(self): grid = Grid(shape=(10, 10)) p0 = TimeFunction(name='p0', grid=grid, space_order=2) f = Function(name='f', grid=grid, space_order=2) eq = Eq(p0.forward, (sin(f)*sin(f)*p0.dx2 + sin(f)*sin(f)*p0.dy2)) op = Operator(eq, opt=opt4()) # Check generated code assert 'pow' not in str(op) @pytest.mark.parametrize('exprs', [ (['u.dz2 + u.dy.dz', 'v.dz2 + v.dy.dz']), (['u.dy.dz + u.dx.dz', 'v.dy.dz + v.dx.dz']), (['u.dx2 + u.dy2 + u.dz2 + u.dx.dy + u.dy.dz + u.dx.dz', 'v.dx2 + v.dy2 + v.dz2 + v.dx.dy + v.dy.dz + v.dx.dz']), ]) @pytest.mark.parametrize('topology', [ (1, '*', 1), (1, 1, '*'), ]) @pytest.mark.parallel(mode=[(4, 'diag2')]) def test_mpi_diag2(self, exprs, topology, mode): grid = Grid(shape=(31, 64, 25), topology=topology) nt = 8 u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) v = TimeFunction(name='v', grid=grid, space_order=4, save=nt, layers=NoLayers) eqns = [] for expr, field in zip(as_tuple(exprs), [u, v]): eqns.append(Eq(field.forward, eval(expr) + 1.)) op0 = Operator(eqns, opt=opt0({'gpu-fit': [u, v], 'cire-mode': '1pass'})) op1 = Operator(eqns, opt=opt4({'par-tile': (10, 7), 'gpu-fit': [u, v], 'cire-mode': '1pass'})) # Check generated code assert 'haloupdate0' in op1._func_table assert 'halowait0' in op1._func_table assert 'haloupdate1' not in op1._func_table assert 'halowait1' not in op1._func_table assert 'gather1' not in op1._func_table init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-2) @pytest.mark.parallel(mode=[(1, 'diag2')]) def test_mpi_custom_partile(self, mode): grid = Grid(shape=(100, 64, 25)) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dy + u.dx + 1.), Eq(v.forward, v.dy + v.dx + 1.)] op = Operator(eqns, opt=opt4({'par-tile': (64, 4, 64)})) # Make sure all symbols are resolved op.cfunction op.arguments(time_M=3) # Check generated code assert 'haloupdate0' in op._func_table assert 'halowait0' in op._func_table assert 'haloupdate1' not in op._func_table assert 'halowait1' not in op._func_table assert 'gather1' not in op._func_table def test_small_halo(self): grid = Grid(shape=(8, 8, 8)) so = 4 u = TimeFunction(name='u', grid=grid, space_order=(so, so // 2, so // 2)) eqn = Eq(u.forward, u.dy.dz + 1.) op = Operator(eqn, opt=opt4({'par-tile': (4, 4, 4)})) op.cfunction @pytest.mark.parametrize('exprs,exp_shmas,exp_queues,exp_syncs,exp_ngas,exp_tmps,kid,oo', [ # noqa (['(p0.dx(x0=x+dx/2) + p0.dy(x0=y-dy/2)).dx(x0=x-dx/2)'], 1, 2, 4, 7, 0, 0, {}), (['(p0.dx(x0=x+dx/2).dx(x0=x-dx/2) + p0.dx(x0=x+dx/2).dy(x0=y-dy/2))'], 1, 1, 2, 3, 0, 0, {}), (['g1t(p0.dx(x0=x+dx/2))'], 2, 2, 2, 3, 0, 1, {'cire-mode': '1pass'}), (['g1t(f * g1(p0))'], 3, 7, 4, 11, 0, 1, {'gpu-loader-threads-radius': (2, 1), 'cire-mode': '1pass'}), (['g1t(f * g1(p0)) + g2t(a * g2(p0))'], 3, 9, 4, 11, 0, 1, {'gpu-loader-threads-radius': (2, 1), 'cire-mode': '1pass'}), (['g1t(f * g1(p0)) + g2t(a * g2(p0))'], (1, 2), (1, 1), (2, 2), (6, 1), (3, 0), (1, 2), {'gpu-loader-threads-radius': (2, 1)}), (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0))'], 3, 10, 4, 11, 0, 1, {'gpu-loader-threads-radius': (2, 1), 'cire-mode': '1pass'}), (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0) + g * g3(q0))', 'g1t(f * g1(q0)) + g2t(a * g2(q0)) + g3t(b * g3(q0) + g * g3(p0))'], 5, 14, 8, 11, 0, 1, {'gpu-loader-threads-radius': (2, 1), 'cire-mode': '1pass'}), # w/ gpu-shm=max-reuse (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0) + g * g3(q0))', 'g1t(f * g1(q0)) + g2t(a * g2(q0)) + g3t(b * g3(q0) + g * g3(p0))'], 4, 14, 6, 11, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads-radius': (2, 1), 'cire-mode': '1pass'}), # w/ gpu-loader-threads: 1D / 1.5D (['g1t(p0.dx(x0=x+dx/2))'], 2, 2, 2, 5, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1D', 'cire-mode': '1pass'}), (['p0.dy(x0=y+dy/2).dy(x0=y-dy/2) + p0.dy(x0=y+dy/2).dz(x0=z-dz/2)'], 2, 0, 3, 4, 0, 0, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1.5D'}), (['p0.dy(x0=y+dy/2).dy(x0=y-dy/2) + p0.dy(x0=y+dy/2).dz(x0=z-dz/2) + ' 'p0.dy(x0=y+dy/2).dx(x0=x-dx/2)'], 2, 2, 4, 7, 0, 0, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1.5D'}), (['(sin(h)*p0.dy(x0=y+dy/2)).dy(x0=y-dy/2) - ' '(cos(h)*p0.dy(x0=y+dy/2)).dx(x0=x-dx/2)'], 1, 3, 8, 3, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1.5D'}), (['g1t(f * g1(p0))'], 2, 7, 6, 7, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1.5D', 'cire-mode': '1pass'}), (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0))'], 2, 10, 6, 7, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1.5D', 'cire-mode': '1pass'}), (['(p0.dx(x0=x+dx/2).dx(x0=x-dx/2) + p0.dx(x0=x+dx/2).dy(x0=y-dy/2))'], 1, 1, 2, 3, 0, 0, None), (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0))'], 3, 10, 4, 3, 0, 1, {'cire-mode': '1pass'}), # Expected to emit "Suboptimal code ..." because of `gpu-loader-threads=1D` (['g1t(f * g1(p0))'], 8, 7, 6, 13, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1D', 'cire-mode': '1pass'}), (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0))'], 8, 10, 7, 13, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1D', 'cire-mode': '1pass'}), # Challenging cases (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0) + g * g3(q0))', 'g1t(f * g1(q0)) + g2t(a * g2(q0)) + g3t(b * g3(q0) + g * g3(p0))'], 4, 14, 6, 7, 0, 1, {'gpu-shm': 'max-reuse', 'gpu-loader-threads': '1.5D', 'cire-mode': '1pass'}), (['g1t(f * g1(p0)) + g2t(a * g2(p0)) + g3t(b * g3(p0) + g * g3(q0))', 'g1t(f * g1(q0)) + g2t(a * g2(q0)) + g3t(b * g3(q0) + g * g3(p0))'], (2, 4), (2, 2), (2, 2), (6, 1), (6, 0), (1, 2), None), ]) def test_tti_sa_blend(self, exprs, exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps, kid, oo): grid = Grid(shape=(31, 17, 25)) nt = 5 x, y, z = grid.dimensions dx, dy, dz = x.spacing, y.spacing, z.spacing # noqa a = Function(name='a', grid=grid, space_order=4) b = Function(name='b', grid=grid, space_order=4) f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) e = Function(name='e', grid=grid, space_order=4) h = Function(name='h', grid=grid, space_order=4) p0 = TimeFunction(name='p0', grid=grid, space_order=4, save=nt, layers=NoLayers) p1 = TimeFunction(name='p0', grid=grid, space_order=4, save=nt, layers=NoLayers) q0 = TimeFunction(name='q0', grid=grid, space_order=4, save=nt, layers=NoLayers) a.data[:] = .05 b.data[:] = .035 f.data[:] = .01 g.data[:] = .017 e.data[:] = .02 h.data[:] = .04 def g1(field): return (cos(e) * cos(h) * field.dx(x0=x+x.spacing/2) + cos(e) * sin(h) * field.dy(x0=y+y.spacing/2) - sin(e) * field.dz(x0=z+z.spacing/2)) def g2(field): return - (sin(h) * field.dx(x0=x+x.spacing/2) - cos(h) * field.dy(x0=y+y.spacing/2)) def g3(field): return (sin(e) * cos(h) * field.dx(x0=x+x.spacing/2) + sin(e) * sin(h) * field.dy(x0=y+y.spacing/2) + cos(e) * field.dz(x0=z+z.spacing/2)) def g1t(field): return ((cos(e) * cos(h) * field).dx(x0=x-x.spacing/2) + (cos(e) * sin(h) * field).dy(x0=y-y.spacing/2) - (sin(e) * field).dz(x0=z-z.spacing/2)) def g2t(field): return - ((sin(h) * field).dx(x0=x-x.spacing/2) - (cos(h) * field).dy(x0=y-y.spacing/2)) def g3t(field): return ((sin(e) * cos(h) * field).dx(x0=x-x.spacing/2) + (sin(e) * sin(h) * field).dy(x0=y-y.spacing/2) + (cos(e) * field).dz(x0=z-z.spacing/2)) exprs = as_tuple(exprs) eqns = [] fields = [] for expr, field in zip(exprs, [p0, q0]): eqns.append(Eq(field.forward, eval(expr) + 1.)) fields.append(field) op0 = Operator(eqns, opt=opt0({'gpu-fit': fields})) op1 = Operator(eqns, opt=opt4({'par-tile': (10, 8), 'gpu-fit': fields, 'deriv-collect': True, # Legacy behavior **(oo or {})})) # Check generated code # NOTE: we force the compiler to generate 1pass implementations if isinstance(kid, int): assert ('kernel%d' % kid) in op1._func_table assert ('kernel%d' % (kid+1)) not in op1._func_table self._check_kernel_opt(op1, p0, exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps, kids=kid) init_chessboard(p0) init_chessboard(p1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, p0=p1) assert np.allclose(p0.data, p1.data, rtol=1.e-4) @pytest.mark.parametrize('exprs,exp_shmas,exp_queues,exp_syncs,exp_ngas,exp_tmps,oo', [ # noqa (['p0.dx(fd_order=so1).dy(fd_order=so1).T'], 1, 1, 2, 3, 0, {}), (['g3t(g3(p0))'], (0, 1, 2), (0, 1, 3), (0, 2, 4), (0, 6, 1), (3, 1, 0), {'cire-mode': '2pass'}), ]) def test_tti_zhang_blend(self, exprs, exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps, oo): grid = Grid(shape=(31, 17, 25)) nt = 5 x, y, z = grid.dimensions dx, dy, dz = x.spacing, y.spacing, z.spacing # noqa so = 8 so1 = so // 2 # noqa e = Function(name='e', grid=grid, space_order=so) h = Function(name='h', grid=grid, space_order=so) p0 = TimeFunction(name='p0', grid=grid, space_order=so, save=nt, layers=NoLayers) p1 = TimeFunction(name='p0', grid=grid, space_order=so, save=nt, layers=NoLayers) q0 = TimeFunction(name='q0', grid=grid, space_order=so, save=nt, layers=NoLayers) e.data[:] = .02 h.data[:] = .04 def g3(field): return (sin(e) * cos(h) * field.dx(fd_order=so1) + sin(e) * sin(h) * field.dy(fd_order=so1) + cos(e) * field.dz(fd_order=so1)) def g3t(field): return ((sin(e) * cos(h) * field).dx(fd_order=so1).T + (sin(e) * sin(h) * field).dy(fd_order=so1).T + (cos(e) * field).dz(fd_order=so1).T) exprs = as_tuple(exprs) eqns = [] fields = [] for expr, field in zip(exprs, [p0, q0]): eqns.append(Eq(field.forward, eval(expr) + 1.)) fields.append(field) op0 = Operator(eqns, opt=opt0({'gpu-fit': fields})) op1 = Operator(eqns, opt=opt4({'par-tile': (10, 8), 'gpu-fit': fields, **(oo or {})})) # Check generated code self._check_kernel_opt(op1, p0, exp_shmas, exp_queues, exp_syncs, exp_ngas, exp_tmps) init_chessboard(p0) init_chessboard(p1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, p0=p1) assert np.allclose(p0.data, p1.data, rtol=1.e-4) def test_perfect_derivs_but_not_just_indexeds(self): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid, space_order=8) g = Function(name='g', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, space_order=8) # These are all perfect derivatives (some mixed, some first order, # some second order), so despite there's a lot of them, the compiler # should still opt for one-pass eq = Eq(u.forward, ((1/u).dx + u.dy.dz + f*u.dz.dx + (1/u).dy + f*u.dx.dy2 + f*u.dz.dy + f*u.dx.dx + (f/u).dz2 + g*u.dx.dy + g*(f*u).dy.dx + f*(1/g*u).dx.dx) + f.dx2 + g.dy2) op0 = Operator(eq, opt=opt4({'cire-hoist': 0, # Legacy 'cire-mode': '1pass'})) op1 = Operator(eq, opt=opt4({'cire-hoist': 0, # Legacy 'gpu-flip-derivs': True})) for op in [op0, op1]: op.cfunction # Check generated code symbols = FindSymbols().visit(op) arrays = [i for i in symbols if isinstance(i, TempArray)] assert len(arrays) == 0 assert 'kernel0' in op._func_table assert 'kernel1' not in op._func_table @pytest.mark.parametrize('exprs,exp_shmas,exp_queues,exp_syncs,exp_ngas', [ # NOTE: different from test_3D's # `('u.dy.dz + u.dx.dz + u.dy.dx', 1, 2, 6, 7, False),` # because the above one factorizes the outer .dz in advance ('u.dy.dz + u.dz.dx + u.dx.dy', 2, 2, 6, 7), # Akin to a subset of terms in the first tau equation in the staggered # TTI elastic formulation ('u.dx + f*u.dy.dx', 1, 2, 4, 7), # no-op ('u.dx + f*u.dx.dy', 1, 2, 4, 7), # Akin to a subset of terms in the first tau equation in the staggered # TTI elastic formulation ('u.dx + u.dx.dz + u.dx.dy', 2, 2, 4, 7), # Akin to the first tau equation in the staggered TTI elastic # formulation ('u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 3, 4, 8, 7), # Simple but with a third order cross derivative # Akin to a subset of terms in the second tau equation in the staggered # TTI elastic formulation ('u.dy.dx + f*w.dy.dz.dx', 1, 2, 10, 7), # Akin to the second tau equation in the staggered TTI elastic # formulation ('u.dy.dx + f*w.dy.dz.dx + u.dy.dz + g*v.dx + u.dy', 2, 3, 7, 7), # Akin to subsets of terms in the first two tau equations in the # staggered TTI elastic formulation ('u.dx.dy + g*v.dy.dx + f*u.dy.dx', 1, 2, 8, 7), ('u.dx.dy + g*v.dy.dx + f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz', 2, 3, 10, 7), # Akin to the first two tau equations in the staggered TTI elastic # formulation (('u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy'), 4, 6, 10, 7), # Akin to the third tau equation in the staggered TTI elastic ('u.dz.dx + v.dy.dz.dx + u.dz.dy + w.dx + u.dz', 3, 4, 11, 7), # Akin to the second and third tau equations in the staggered TTI elastic # formulation (('f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy', 'u.dz.dx + v.dy.dz.dx + u.dz.dy + w.dx + u.dz'), 2, 6, 12, 7), # Akin to subsets of terms in the first and third tau equations in the # staggered TTI elastic formulation (('u.dx.dy', 'u.dz.dy'), 1, 1, 6, 7), (('u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'u.dz.dx + f*v.dy.dz.dx + u.dz.dy + g*w.dx + u.dz'), 5, 4, 7, 7), # Akin to the first three tau equations in the staggered TTI elastic # formulation (('u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy', 'u.dz.dx + f*v.dy.dz.dx + u.dz.dy + g*w.dx + u.dz'), 6, 6, 6, 7), # Akin to the first four tau equations in the staggered TTI elastic # formulation (('u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy', 'u.dz.dx + f*v.dy.dz.dx + u.dz.dy + g*w.dx + u.dz', 'v.dy + f*u.dx.dy + v.dy.dz + v.dy.dx + g*w.dx.dy'), 7, 6, 6, 7), # Akin to the first five tau equations in the staggered TTI elastic # formulation (('u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy', 'u.dz.dx + f*v.dy.dz.dx + u.dz.dy + g*w.dx + u.dz', 'v.dy + f*u.dx.dy + v.dy.dz + v.dy.dx + g*w.dx.dy', 'v.dz.dy + f*u.dx.dz.dy + v.dz.dx + g*w.dy + v.dz'), 10, 6, 7, 7), # Akin to the staggered TTI elastic formulation (('u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy', 'u.dz.dx + f*v.dy.dz.dx + u.dz.dy + g*w.dx + u.dz', 'v.dy + f*u.dx.dy + v.dy.dz + v.dy.dx + g*w.dx.dy', 'v.dz.dy + f*u.dx.dz.dy + v.dz.dx + g*w.dy + v.dz', 'w.dz + f*u.dx.dz + g*v.dy.dz + w.dz.dy + w.dz.dx'), 12, 6, 6, 7), # Lifted from test_3D_lt1D (('u.dy.dz + f*u.dz.dy',), 2, 0, 3, 4), ]) @switchconfig(ignore_unknowns=True) def test_flip(self, exprs, exp_shmas, exp_queues, exp_syncs, exp_ngas): grid = Grid(shape=(31, 17, 25)) nt = 5 # These are necessary to prevent derivative factorization f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) functions = [] for name in ('u', 'v', 'w', 'a', 'b', 'c'): f0 = TimeFunction(name=name, grid=grid, space_order=4, save=nt, layers=NoLayers) f1 = f0.func(name='%s1' % name) functions.extend([f0, f1]) u, u1, v, v1, w, w1, a, a1, b, b1, c, c1 = functions f.data[:] = 19.0 g.data[:] = 18.0 exprs = as_tuple(exprs) fields = [u, v, w, a, b, c] eqns = [] for expr, field in zip(exprs, fields): eqns.append(Eq(field.forward, eval(expr) + 1.)) op0 = Operator(eqns, opt=opt0({'gpu-fit': fields})) op1 = Operator(eqns, opt=opt4({'par-tile': (10, 7), 'gpu-fit': fields, 'gpu-flip-derivs': True, 'deriv-collect': True, # Legacy behavior 'cire-mode': '1pass'})) # Check generated code # NOTE: these are expected to be all 1-pass implementations assert 'kernel0' in op1._func_table assert 'kernel1' not in op1._func_table self._check_kernel_opt(op1, u, exp_shmas, exp_queues, exp_syncs, exp_ngas) for i in (u, u1, v, v1, w, w1): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1, v=v1, w=w1) assert np.allclose(u.data, u1.data, rtol=1.e-4) @pytest.mark.parametrize('orders', [(0, 1), (1, 0), (1, 2), (0, 1, 2)]) def test_flip_deriv_order(self, orders): grid = Grid(shape=(11,)*len(orders)) dims = grid.dimensions f = Function(name='f', grid=grid, space_order=4) expr = f for d, order in zip(dims, orders): x0 = d + d.spacing/2 if order == 0 else d expr = expr.diff(d, deriv_order=order, fd_order=4, x0=x0) ideriv = diff2sympy(expr._evaluate(expand=False)) # Reverse the nesting while retaining each axis' derivative order priority = {d: -i for i, d in enumerate(dims)} subs, flag = _flip([ideriv], priority) assert not flag flipped = subs[ideriv] for d, order in zip(dims, orders): assert tuple(flipped.mapper) == (d,) assert flipped.deriv_order == order flipped = flipped.base @pytest.mark.parametrize('exprs,exp_shmas,exp_queues,exp_syncs,exp_ngas,extra,relax_if_amd', [ # noqa # NOTE: some tests are skipped on AMD as they would crash for totally # unexplainable reasons; same test runs perfectly fine under nvidia # debugger (so no memory violations); autopadding=0 no change; removing # any random line from the generated code with the jit backdoor makes # the test pass; any space order lower than 16 makes the test pass; # more complicated tests below do pass # Akin to subsets of terms in the first two tau equations in the staggered # TTI elastic formulation (('vy.dx + tau[0]', 'vy.dx + tau[1]'), 1, 2, 4, 0, {}, False), (('vx.dx + vx.dz + tau[0]', 'tau[1]'), 1, 1, 2, 2, {}, False), (('vx.dx + vx.dz + tau[0]', 'vx.dx'), 2, 2, 4, 4, {}, False), (('vx.dx + vx.dz + tau[0]', 'vx.dx + vx.dy + tau[1]'), 2, 3, 4, 6, {}, False), # MFE for what was a missing __syncthreads() in the generated code (('vx.dz + tau[0]', 'tau[1]', 'f*vy.dx + tau[2]'), 3, 3, 6, 2, {}, True), # Inspired by the first two tau equations in the staggered TTI elastic pytest.param( ('vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + tau[0]', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + tau[1]'), 7, 12, 10, 8, {}, 1.e-2, marks=skipif('device-hip') ), # Inspired by the first three tau equations in the staggered TTI elastic pytest.param( ('vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + tau[0]', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + tau[1]', 'vx.dx + f*vy.dx + vx.dy + g*vz.dx + vx.dz + tau[2]'), 11, 12, 10, 8, {}, 1.e-2, marks=skipif('device-hip') ), # Odd MFE for code gen issue we had, inspired by the first four tau # equations in the staggered TTI elastic pytest.param( ('3. + .01*sum(fields)', '1. + .02*sum(fields)', '0.02*vy.dx + .2 + .03*sum(fields)', '0.01*vz.dy + .1 + .04*sum(fields)',), 13, 11, 9, 0, {}, False, marks=skipif('device-hip') ), # Inspired by staggered TTI elastic pytest.param( ('vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + .001*sum(fields)', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + .002*sum(fields)', '0.1*vx.dx + f*vy.dx + vx.dy + g*vz.dx + 0.4*vx.dz + .003*sum(fields)', 'vy.dy + f*vx.dy + vy.dz + vy.dx + g*vz.dy + .004*sum(fields)', 'vy.dy + f*vx.dy + vy.dx + g*vz.dy + vy.dz + .005*sum(fields)', 'vz.dz + f*vx.dz + g*vy.dz + vz.dy + vz.dx + .006*sum(fields)'), 34, 28, 23, 7, {}, 1.e-1, marks=skipif('device-hip') ), (('f*(vx.dz+vz.dx)', 'g*(vx.dz+vz.dx)'), 5, 4, 7, 3, {}, True), pytest.param( ('vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)',), 6, 5, 7, 6, {}, 1.e-2, marks=skipif('device-hip') ), (('.01', 'vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)'), 8, 8, 10, 6, {}, 1.e-2), pytest.param( ('vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)', 'vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)'), 13, 16, 15, 8, {}, False, marks=skipif('device-hip') ), (('vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)', 'vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)', 'vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)', 'vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)', 'vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)', 'vx.dx+vy.dy+vz.dz+f*(vx.dy+vy.dx)+g*(vx.dz+vz.dx)+f*(vy.dz+vz.dy)'), 30, 16, 11, 7, {}, 1.e-1), # With `gpu-loader-threads` disabled (('.01', '.02', 'g*vy.dz + .03'), 2, 2, 6, 0, {'gpu-loader-threads': False}, False), (('vx.dx + tau[0]', 'vx.dz + tau[1]'), 2, 1, 3, 7, {'gpu-loader-threads': False}, False), (('vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + tau[2]',), 5, 5, 4, 10, {'gpu-loader-threads': False, 'par-tile': (16, 16)}, False), (('vx.dy + tau[0]', 'vx.dy + tau[1]'), 1, 1, 4, 6, {'gpu-loader-threads': False, 'par-tile': (16, 16)}, False), (('tau[0]', 'vx.dz + tau[1]'), 2, 0, 3, 6, {'gpu-loader-threads': False}, False), (('vx.dz + vx.dy + tau[0]', 'vx.dz + vx.dy + tau[1]'), 3, 2, 5, 14, {'gpu-loader-threads': False, 'par-tile': (16, 16)}, False), pytest.param( ('vx.dx + vx.dz + vx.dy + tau[2]', 'vx.dx + vx.dz + vx.dy + tau[1]', 'vx.dx + vx.dz + vx.dy + tau[0]'), 8, 7, 5, 32, {'gpu-loader-threads': False, 'par-tile': (16, 16)}, False, marks=skipif('device-hip') ), ]) @switchconfig(ignore_unknowns=True) def test_flip_staggering(self, exprs, exp_shmas, exp_queues, exp_syncs, exp_ngas, extra, relax_if_amd): grid = Grid(shape=(31, 17, 25)) x, y, z = grid.dimensions nt = 5 # These are necessary to prevent derivative factorization f = Function(name='f', grid=grid, space_order=16) g = Function(name='g', grid=grid, space_order=16) v = VectorTimeFunction(name="v", grid=grid, space_order=16) tau = TensorTimeFunction(name="tau", grid=grid, space_order=16, save=nt, layers=NoLayers) fields = tau.values() vx, vy, vz = v # noqa f.data[:] = 19.0 g.data[:] = 18.0 v[0].data[:] = 11.0 v[1].data[:] = 12.0 v[2].data[:] = 7.0 exprs = as_tuple(exprs) eqns = [] for expr, field in zip(exprs, fields): eqns.append(Eq(field.forward, eval(expr) + 1.)) op0 = Operator(eqns, opt=opt0({'gpu-fit': fields, 'cire-mode': '1pass', 'fission': 'aggressive:300'})) op1 = Operator(eqns, opt=opt4({**{'errctl': 'basic', 'par-tile': (16, 8), 'gpu-fit': fields, 'gpu-flip-derivs': True, 'deriv-collect': True, # Legacy 'cire-hoist': 0, # Legacy 'cire-mode': '1pass', 'fission': 'aggressive:300'}, **extra})) # Check generated code # NOTE: these are expected to be all 1-pass implementations assert 'kernel0' in op1._func_table assert 'kernel1' not in op1._func_table self._check_kernel_opt(op1, v[0], exp_shmas, exp_queues, exp_syncs, exp_ngas) overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} for objs in (fields, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) # AMD has a tiny larger more significant loss of accuracy in some tests if relax_if_amd and isinstance(configuration['platform'], AmdDevice): rtol = relax_if_amd else: rtol = 1.e-3 for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=rtol, atol=1e-3) def test_flip_for_canonical_ordering(self): grid = Grid(shape=(31, 17, 25)) nt = 5 f = Function(name='f', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) v = TimeFunction(name='v', grid=grid, space_order=4) eqn = Eq(u.forward, v.dx.dy + (f*v.dy).dz.dx + (v.dx.dy + f).dz + u + 1) op0 = Operator(eqn, opt=opt0()) op1 = Operator(eqn, opt=opt4({'par-tile': (10, 8), 'gpu-flip-derivs': True})) # Check generated code self._check_kernel_opt(op1, v, 2, 2, 6, 6) launches = FindNodes(KernelLaunch).visit(op1) assert launches[0].block.cargs == (14, 8) # Loader threads only along z u1 = u.func(name='u1') f.data[:] = .0001 init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) def test_flip_and_2pass_for_canonical_ordering(self): grid = Grid(shape=(31, 17, 25)) nt = 5 f = Function(name='f', grid=grid, space_order=4) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) v = TimeFunction(name='v', grid=grid, space_order=4) eqn = Eq(u.forward, v.dy.dx + (f*v.dx).dz.dy + g*(v.dz.dx + f).dy + u + 1) op0 = Operator(eqn, opt=opt0()) op1 = Operator(eqn, opt=opt4({'par-tile': (10, 8), 'gpu-flip-derivs': True})) # Check generated code self._check_kernel_opt(op1, v, (1, 3), (1, 1), (2, 6), (2, 6), exp_tmps=(2, 0), kids=(0, 1)) launches = FindNodes(KernelLaunch).visit(op1) assert len(launches) == 2 assert launches[0].block.cargs == (14, 8) # Loader threads only along z assert launches[1].block.cargs == (14, 8) # Loader threads only along z u1 = u.func(name='u1') f.data[:] = .035 g.data[:] = .021 init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) @pytest.mark.parametrize('expr', [ ('u.dx + f*u.dy.dx'), ('u.dx2 + u.dy2 + u.dz2 + u.dx.dy + u.dy.dz + f*u.dx.dz'), ]) def test_partile_smaller_than_halo(self, expr): grid = Grid(shape=(100, 64, 25)) nt = 5 f = Function(name='f', grid=grid, space_order=16) u = TimeFunction(name='u', grid=grid, space_order=16, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=16, save=nt, layers=NoLayers) eqns = [Eq(u.forward, eval(expr) + 1.)] op0 = Operator(eqns, opt=opt0({'gpu-fit': [u]})) op1 = Operator(eqns, opt=opt4({'par-tile': (32, 4, 16), 'gpu-fit': [u]})) f.data[:] = .0001 init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def test_dy_of_preeval_deriv(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4) eq = Eq(u.forward, (u.dx.evaluate).dy + 1) op = Operator(eq, opt=opt4({'par-tile': (4, 4)})) op.cfunction def test_kernel_guard_w_cond_dim(self): grid = Grid(shape=(10, 10, 10)) time_dim = grid.time_dim factor = Constant(name='factor', value=2, dtype=np.int32) cd = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor) f = Function(name='f', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4) eq = Eq(f, u.laplace + 1, implicit_dims=cd) op = Operator(eq, opt=opt4({'par-tile': (4, 4, 4)})) assert len(op._func_table) == 7 assert 'kernel0' in op._func_table assert 'kernel1' not in op._func_table @pytest.mark.parametrize('opt,expected', [ (opt4({'gpu-opt': True}), True), # Below, the threshold is instead set programmatically, so the outcome # should be the same for both CUDA and HIP (opt4({'gpu-unroll-ths': 10}), True), (opt4({'gpu-unroll-ths': 16}), False) ]) def test_unroll_if_necessary(self, opt, expected): grid = Grid(shape=(51, 33, 42)) so = 12 u = TimeFunction(name='u', grid=grid, space_order=so) eqn = Eq(u.forward, u.dx + 1) op = Operator(eqn, opt=opt) i0, i1, i2, i3 = FindNodes(Iteration).visit(op._func_table['kernel0']) if expected: assert str(i0.pragmas[0]) == f'#pragma unroll({so})' assert not i1.pragmas assert str(i2.pragmas[0]) == f'#pragma unroll({so})' assert str(i3.pragmas[0]) == f'#pragma unroll({so+1})' else: assert not any(i.pragmas for i in (i0, i1, i2, i3)) def test_expand_if_tiny(self): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4) eqn = Eq(u.forward, u.laplace + f.dy(fd_order=1).dz(fd_order=1)) op = Operator(eqn, opt=opt4({'cire-hoist': 0, # Legacy 'par-tile': (4, 4, 4)})) op.cfunction # Check generated code -- expansion was performed kernel = op._func_table['kernel0'].root assert 'fL0(x + 4, y + 5, z + 4)' in str(kernel) # ... and no kernel augmentation (i.e., loader threads) took place! launch, = FindNodes(KernelLaunch).visit(op) assert launch.block.cargs == (4, 4) @pytest.mark.parametrize('optoptions,exp_shmas,exp_queues,exp_syncs,exp_ngas', [ ({}, 30, 16, 11, 7), ({'gpu-loader-threads': False}, 31, 19, 10, 33), ]) def test_etti_stress_update_like(self, optoptions, exp_shmas, exp_queues, exp_syncs, exp_ngas): grid = Grid(shape=(512, 512, 512)) so = 8 f = Function(name='f', grid=grid, space_order=so) g = Function(name='g', grid=grid, space_order=so) v = VectorTimeFunction(name="v", grid=grid, space_order=so) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so) vx, vy, vz = v expr = (vx.dx + vy.dy + vz.dz + f*(vx.dy + vy.dx) + g*(vx.dz + vz.dx) + f*(vy.dz + vz.dy)) eqns = [Eq(i.forward, expr) for i in tau.values()] op = Operator(eqns, opt=opt4({**{'cire-hoist': 0, # Legacy 'par-tile': (24, 4, 32), 'gpu-flip-derivs': True}, **optoptions})) # Check generated code assert 'kernel0' in op._func_table assert 'kernel1' not in op._func_table self._check_kernel_opt(op, v[0], exp_shmas, exp_queues, exp_syncs, exp_ngas) op.apply(time_M=0) def test_issue_707_v0(self): grid = Grid(shape=(11, 11, 11)) x, y, z = grid.dimensions so = 8 x0p, y0p, z0p = x + x.spacing/2, y + y.spacing/2, z + z.spacing/2 x0m, y0m, z0m = x - x.spacing/2, y - y.spacing/2, z - z.spacing/2 theta = Function(name='theta', grid=grid, space_order=so) phi = Function(name='phi', grid=grid, space_order=so) ct, st = cos(theta), sin(theta) cp, sp = cos(phi), sin(phi) s2p = sin(2 * phi) R = np.array([[1 - (1 - ct) * cp**2, -1/2 * s2p * (1 - ct), -st * cp], [-1/2 * s2p * (1 - ct), 1 - (1 - ct) * sp**2, -st * sp], [st * cp, st * sp, ct]]) cijs = [Function(name='c%d%d' % (i+1, j+1), grid=grid, space_order=so) for i in range(3) for j in range(3)] v1 = TimeFunction(name='v1', grid=grid, space_order=so, time_order=1, save=Buffer(1)) tau = TimeFunction(name='tau', grid=grid, space_order=so, time_order=1, save=Buffer(1)) tau2 = TimeFunction(name='tau2', grid=grid, space_order=so, time_order=1, save=Buffer(1)) KW = dict(deriv_order=(0, 0), fd_order=(2, 2)) xy_m_p = {x: x0m, y: y0p} xy_p_m = {x: x0p, y: y0m} xz_m_p = {x: x0m, z: z0p} yz_m_p = {y: y0m, z: z0p} def dx(field): return ((R[0, 0]*field).dx(x0=x0m) + Derivative(R[0, 1]*field, x, y, x0=xy_m_p, **KW).dy(x0=y0m) + Derivative(R[0, 2]*field, x, z, x0=xz_m_p, **KW).dz(x0=z0m)) def dy(field): return (Derivative(R[1, 0]*field, x, y, x0=xy_p_m, **KW).dx(x0=x0m) + (R[1, 1]*field).dy(x0=y0m) + Derivative(R[1, 2]*field, y, z, x0=yz_m_p, **KW).dz(x0=z0m)) v1f = v1.forward eqns = [Eq(v1.forward, v1 - div(tau)), Eq(tau.forward, tau - cijs[0]*dx(cos(theta)*v1f) - cijs[1]*dy(sin(theta)*v1f)), Eq(tau2.forward, tau2 - cijs[3]*dx(cos(theta)*v1f) - cijs[4]*dy(sin(theta)*v1f))] op = Operator(eqns, opt=opt4({'errctl': 'basic', 'gpu-opt': True, 'gpu-flip-derivs': True, 'par-tile': (8, 8, 8)})) # Ensure it runs to completion w/o errors op(time=1) # Check generated code self._check_kernel_opt(op, tau, 4, 5, 10, 1, kids=2) @switchconfig(autopadding=False) # To expose hazardous OOB accesses def test_issue_707_v1(self): space_order = 8 interp_order = 2 kw = dict(deriv_order=(0, 0), fd_order=(interp_order, interp_order)) grid = Grid((131, 101, 150)) x, y, z = grid.dimensions t = grid.stepping_dim x0m, y0m, z0m = x - x.spacing/2, y - y.spacing/2, z - z.spacing/2 x0p, y0p, z0p = x + x.spacing/2, y + y.spacing/2, z + z.spacing/2 f = Function(name='f', grid=grid, space_order=space_order) g = Function(name='g', grid=grid, space_order=space_order) D = Derivative def foo0(field): return (D(sin(f)*field, x, y, x0={x: x0p, y: y0m}, **kw).dx(x0=x0m) + (cos(g)*field).dy(x0=y0m) + D(sin(g)*field, y, z, x0={y: y0m, z: z0p}, **kw).dz(x0=z0m)) def foo1(field): return (D(cos(f)*field, x, z, x0={x: x0p, z: z0m}, **kw).dx(x0=x0m) + D(sin(g)*field, y, z, x0={y: y0p, z: z0m}, **kw).dy(x0=y0m) + (cos(g)*field).dz(x0=z0m)) def foo2(field): return ((cos(f)*field).dx(x0=x0p) + D(sin(g)*field, x, z, x0={x: x0p, z: z0m}, **kw).dz(x0=z0p)) def foo3(field): return (D(field, x, z, x0={x: x0m, z: z0p}, **kw).dx(x0=x0p) + field.dz(x0=z0p)) kwargs = {'grid': grid, 'space_order': space_order, 'time_order': 1, 'save': Buffer(1), 'is_transient': False} vx = TimeFunction(name='vx', **kwargs) vy = TimeFunction(name='vy', **kwargs) vz = TimeFunction(name='vz', **kwargs) txx = TimeFunction(name='txx', **kwargs) txz = TimeFunction(name='txz', **kwargs) expr0 = (foo0(vy.forward) + foo1(vz.forward)) expr2 = (foo2(vz.forward) + foo3(vx.forward)) eqns = [Eq(txx.forward, t.spacing * expr0 + txx + .1), Eq(txz.forward, t.spacing * expr2 + txz + .2)] subs = grid.spacing_map subs[grid.time_dim.spacing] = 1 op0 = Operator(eqns, subs=subs, opt=opt0()) op1 = Operator(eqns, subs=subs, opt=opt4({'errctl': 'basic'})) txx1 = txx.func(name='txx1') txz1 = txz.func(name='txz1') # NOTE: We pollute the halo with NaNs to catch any hazardous # out-of-bounds accesses for i in (f, g, txx, txx1, txz, txz1): init_chessboard(i) op0(dt=0.001, time=3) op1(dt=0.001, time=3, txx=txx1, txz=txz1) assert np.allclose(txx.data, txx1.data, rtol=1.e-5) assert np.allclose(txz.data, txz1.data, rtol=1.e-5) def test_issue_707_v2(self): grid = Grid(shape=(31, 17, 25)) x, y, z = grid.dimensions nt = 5 f = Function(name='f', grid=grid, space_order=8) v = VectorTimeFunction(name="v", grid=grid, space_order=8) tau = TensorTimeFunction(name="tau", grid=grid, space_order=8, save=nt, layers=NoLayers) fields = tau[:3] vx, vy, vz = v # noqa f.data[:] = 19.0 v[0].data[:] = 11.0 v[1].data[:] = 12.0 v[2].data[:] = 7.0 eqns = [Eq(tau[0].forward, vx.dx + vy.dy + f*vy.dx + tau[0] + .1), Eq(tau[1].forward, vx.dx + vy.dy + vz.dz + tau[1] + .1), Eq(tau[2].forward, vy.dx + tau[2] + .1)] op0 = Operator(eqns, opt=opt0({'gpu-fit': fields})) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': (16, 8, 8), 'gpu-fit': fields, 'gpu-flip-derivs': True, 'cire-schedule': -1})) overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} for objs in (fields, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) def test_fuse_queue_loops(self): grid = Grid(shape=(40, 30, 20)) x, y, z = grid.dimensions space_order = 8 interp_order = 2 vx = TimeFunction(name='vx', grid=grid, space_order=space_order, time_order=1, save=Buffer(1)) txx = vx.func(name='txx') txy = vx.func(name='txy') txz = vx.func(name='txz') f = Function(name='f', grid=grid, space_order=space_order) x0m = x - x.spacing/2 y0m = y - y.spacing/2 z0m = z - z.spacing/2 x0p = x + x.spacing/2 y0p = y + y.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(interp_order, interp_order)) expr0 = (txx.dx(x0=x0p) + Derivative(txx.dy(x0=y0p), x, y, x0={x: x0p, y: y0m}, **kw)) expr1 = (Derivative(txy.dx(x0=x0m), x, y, x0={x: x0p, y: y0m}, **kw) + txy.dy(x0=y0m)) expr2 = Derivative(txz.dx(x0=x0m), x, z, x0={x: x0p, z: z0m}, **kw) eqn = Eq(vx.forward, expr0 + f*expr1 + f*expr2 + vx + .001) op0 = Operator(eqn, subs=grid.spacing_map, opt=opt0()) op1 = Operator(eqn, subs=grid.spacing_map, opt=opt4({'errctl': 'basic', 'par-tile': (16, 8), 'gpu-flip-derivs': True})) # Check generated code self._check_kernel_opt(op1, vx, 3, 5, 4, 1) dims = FindSymbols('dimensions').visit(op1._func_table['kernel0']) assert len([d for d in dims if isinstance(d, QueueInitDimension)]) == 1 vx1 = vx.func(name='vx1') f.data[:] = .0001 for i in (vx, vx1, txx, txy, txz): init_chessboard(i) op0.apply(dt=.001, time_M=3) op1.apply(dt=.001, time_M=3, vx=vx1) assert np.allclose(vx.data, vx1.data, rtol=1.e-3) def test_late_rotation_w_reused_queue_dimension(self): grid = Grid(shape=(20, 16, 12)) x, y, _ = grid.dimensions vx = TimeFunction(name='vx', grid=grid, space_order=4, time_order=1) vx1 = vx.func(name='vx1') txx = vx.func(name='txx') txy = vx.func(name='txy') f = Function(name='f', grid=grid, space_order=4) x0m = x - x.spacing/2 y0m = y - y.spacing/2 x0p = x + x.spacing/2 y0p = y + y.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(2, 2)) expr0 = Derivative( txx.dy(x0=y0p), x, y, x0={x: x0p, y: y0m}, **kw ) expr1 = (Derivative( txy.dx(x0=x0m), x, y, x0={x: x0p, y: y0m}, **kw ) + txy.dy(x0=y0m)) eqn = Eq(vx.forward, expr0 + f*expr1) op0 = Operator(eqn, subs=grid.spacing_map, opt=opt0()) op1 = Operator( eqn, subs=grid.spacing_map, opt=opt4({ 'errctl': 'basic', 'par-tile': (8, 4), 'gpu-flip-derivs': True, }) ) # The eager and late rotations have identical physical bounds but # different logical fronts, so they require distinct Dimensions. kernel = op1._func_table['kernel0'].root queues = [i for i in FindSymbols().visit(kernel) if isinstance(i, TBQueue)] rds = [i.dimensions[0] for i in queues if i.dimensions[0]._size == 2] assert len(rds) == 2 assert len(set(rds)) == 2 assert sorted(i._max_log for i in rds) == [1, 2] f.data[:] = 2 init_chessboard(txx) init_chessboard(txy) op0.apply(time_M=0) op1.apply(time_M=0, vx=vx1) assert np.allclose(vx.data, vx1.data, rtol=1.e-5) def test_issue_947(self): grid = Grid(shape=(20, 16, 12), extent=(19., 15., 11.)) a = Function(name='a', grid=grid, space_order=8) f = Function(name='f', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=(8, 4, 4), save=Buffer(2)) u1 = u.func(name='u1') q = a*u eqn = Eq(u.backward, q.dz2 + q.dx.dz + (f*q).dz2 + (f*q).dx.dz) op0 = Operator(eqn, name='Issue947Ref', subs=grid.spacing_map, opt=opt0({'errctl': 'basic'})) op1 = Operator(eqn, name='Issue947Opt', subs=grid.spacing_map, opt=opt4({'errctl': 'basic', 'gpu-flip-derivs': True, 'par-tile': (8, 4)})) # Check generated code self._check_kernel_opt(op1, u, 4, 2, 3, 3) a.data[:] = 1.8 f.data[:] = .2 init_chessboard(u) init_chessboard(u1) # A single backward iteration is sufficient to expose the corruption. op0.apply(time_m=0, time_M=0) op1.apply(time_m=0, time_M=0, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def test_allocate_queue_with_nonperfect_derivs(self): grid = Grid(shape=(40, 30, 20)) x, y, z = grid.dimensions space_order = 16 interp_order = 2 vx = TimeFunction(name='vx', grid=grid, space_order=space_order, time_order=1, save=Buffer(1)) txx = vx.func(name='txx') f = Function(name='f', grid=grid, space_order=space_order) g = f.func(name='g') c11 = f.func(name='c11') x0m = x - x.spacing/2 y0m = y - y.spacing/2 y0p = y + y.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(interp_order, interp_order)) R00 = 1-(1-cos(g))*cos(f)**2 R01 = -(1-cos(g))*sin(f)*cos(f) expr0 = ((R00*vx).dx(x0=x0m) + Derivative(R01*vx, x, y, x0={x: x0m, y: y0p}, **kw).dy(x0=y0m)) eqn = Eq(txx.forward, c11*expr0 + txx) op0 = Operator(eqn, subs=grid.spacing_map, opt=opt0()) op1 = Operator(eqn, subs=grid.spacing_map, opt=opt4({'errctl': 'basic', 'par-tile': (16, 8), 'cire-mode': '1pass'})) # Check generated code self._check_kernel_opt(op1, vx, 1, 2, 4, 8, kids=1) txx1 = txx.func(name='txx1') c11.data[:] = 12.5 f.data[:] = 0.25 g.data[:] = 0.5 for i in (vx, txx, txx1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, txx=txx1) assert np.allclose(txx.data, txx1.data, rtol=1.e-5) def test_ensure_2pass_if_not_flippable(self): grid = Grid(shape=(40, 30, 20)) x, y, z = grid.dimensions space_order = 8 interp_order = 2 vx = TimeFunction(name='vx', grid=grid, space_order=space_order, time_order=1, save=Buffer(1)) vy = vx.func(name='vy') vz = vx.func(name='vz') txx = vx.func(name='txx') txz = vx.func(name='txz') tyz = vx.func(name='tyz') f = Function(name='f', grid=grid, space_order=space_order) c11 = f.func(name='c11') x0m, z0m = x - x.spacing/2, z - z.spacing/2 x0p, y0p, z0p = x + x.spacing/2, y + y.spacing/2, z + z.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(interp_order, interp_order)) def foo(field): deriv_xz = Derivative(sin(f)*field, x, z, x0={x: x0m, z: z0p}, **kw) return deriv_xz.dx(x0=x0p) expr0 = (cos(f)*vx).dx(x0=x0m) expr1 = (cos(f)*vz).dx(x0=x0p) expr2 = Derivative(cos(f)*vz, y, z, x0={y: y0p, z: z0m}, **kw).dz(x0=z0p) eqns = [ Eq(txx.forward, c11*expr0 + txx), Eq(tyz.forward, foo(vy) + expr2 + tyz), Eq(txz.forward, foo(vx) + expr1 + txz), ] op = Operator(eqns, subs=grid.spacing_map, opt=opt4({'errctl': 'basic', 'gpu-vect': False, 'gpu-flip-derivs': True, 'deriv-collect': True, # Legacy behavior 'par-tile': (16, 8)})) op.cfunction # Check generated code self._check_kernel_opt(op, vx, 1, 1, 2, 2, 1, kids=1) self._check_kernel_opt(op, vx, 3, 4, 6, 2, kids=2) def test_single_front_w_nonatom_derivs_v0(self): grid = Grid(shape=(40, 40, 40)) x, y, z = grid.dimensions f = Function(name='f', grid=grid, space_order=8) vz = TimeFunction(name='vz', grid=grid, space_order=8, save=Buffer(1)) tyz = vz.func(name='tyz') txz = vz.func(name='txz') z0m = z - z.spacing/2 x0p = x + x.spacing/2 y0p = y + y.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(2, 2)) expr0 = ((sin(f)*vz).dy(x0=y0p) + Derivative(sin(f)*vz, y, z, x0={y: y0p, z: z0m}, **kw)) expr1 = ((cos(f)*vz).dx(x0=x0p) + Derivative(vz, x, z, x0={x: x0p, z: z0m}, **kw)) eqns = [Eq(tyz.forward, expr0 + tyz), Eq(txz.forward, expr1 + txz)] op = Operator(eqns, opt=opt4({'par-tile': (24, 4, 48)})) op.cfunction # Check generated code self._check_kernel_opt(op, vz, 3, 2, 3, 4, kids=1) def test_fission_to_avoid_explosion_in_registers(self): grid = Grid(shape=(20, 20, 20)) # Going to need a huge number of registers... funcs = [TimeFunction(name=f'f{i}', grid=grid, space_order=32) for i in range(11)] eqns = [Eq(f.forward, f.dx) for f in funcs] op = Operator(eqns, opt=('advanced', {'errctl': 'basic', 'par-tile': (32, 4), 'gpu-opt': True, # Legacy behavior 'cire-mode': '1pass', 'fission': 'aggressive:200'})) # Check generated code launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == 2 # Ensure it runs to completion op.apply(time_M=9) def test_issue_795(self): grid = Grid(shape=(100, 100, 100)) x, y, z = grid.dimensions so = 2 f0 = Function(name='f0', grid=grid, space_order=so, staggered=NODE) f1 = Function(name='f1', grid=grid, space_order=so, staggered=NODE) p0 = TimeFunction(name='p0', grid=grid, space_order=so, staggered=(x, z)) p1 = TimeFunction(name='p1', grid=grid, space_order=so, staggered=(y, z)) eqns = [Eq(p0, f0), Eq(p1, 2*(f1 + 1/f0) + f0)] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt4({'par-tile': (16, 16, 16)})) p0_1 = p0.func(name='p0_1') p1_1 = p1.func(name='p1_1') f0.data[:] = 3.0 f1.data[:] = 4.0 for i in (p0, p0_1, p1, p1_1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, p0=p0_1, p1=p1_1) assert np.allclose(p0.data, p0_1.data, rtol=1.e-4) assert np.allclose(p1.data, p1_1.data, rtol=1.e-4) @pytest.mark.parametrize('optoptions,expected', [ ({'cire-hoist': False}, False), ({'cire-hoist': 1}, True), ({'cire-hoist': '1:80'}, True), ({'cire-hoist': '1:82'}, False), # Won't hoist as `len(stencil) == 81` ]) def test_cire_hoist_heuristic(self, optoptions, expected): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, space_order=8) eq = Eq(u.forward, u.laplace + u*(f.dx.dy + f.dx.dz)) op = Operator(eq, opt=opt4({'gpu-flip-derivs': True, **optoptions})) op.cfunction # Check generated code launches = FindNodes(KernelLaunch).visit(op) if expected: assert len(launches) == 2 # Check kernel0 is launched outside of the time loop section0 = op.body.body[0].body[0] assert section0.is_Section assert section0.body[0].body[0].body[0] is launches[0] else: assert len(launches) == 1 def test_issue_z_only_function(self): grid = Grid(shape=(41, 41, 47)) x, y, z = grid.dimensions time = grid.time_dim f = Function(name='f', grid=grid, space_order=16) r = Function(name='r_z1', grid=grid, shape=(47,), space_order=16, dimensions=(z,)) p = TensorTimeFunction(name='p', grid=grid, space_order=16, time_order=1, save=Buffer(1)) v = VectorTimeFunction(name='v', grid=grid, space_order=16, time_order=1, save=Buffer(1)) p_xx, p_xy, p_xz, p_yy, p_yz, p_zz = p.values() v_x, v_y, v_z = v.values() s = time.spacing eqns = [ Eq(v_x.forward, (r*p_xz.dz + s*p_xx.dx + s*p_xy.dy) + v_x / (s*f + 1)), Eq(v_y.forward, (r*p_yz.dz + s*p_xy.dx + s*p_yy.dy) + v_y / (s*f + 1)), Eq(v_z.forward, (r*p_zz.dz + s*p_xz.dx + s*p_yz.dy) + v_z / (s*f + 1)) ] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt4({'par-tile': (16, 16, 16)})) v_x1 = v_x.func(name='v_x1') v_y1 = v_y.func(name='v_y1') v_z1 = v_z.func(name='v_z1') for i in (p_xx, p_xy, p_xz, p_yy, p_yz, p_zz): init_chessboard(i) r.data[:] = 1.0 op0(dt=2.0, time=20) op1(dt=2.0, time=20, v_x=v_x1, v_y=v_y1, v_z=v_z1) assert np.allclose(v_x.data, v_x1.data, rtol=1.e-4) assert np.allclose(v_y.data, v_y1.data, rtol=1.e-4) assert np.allclose(v_z.data, v_z1.data, rtol=1.e-4) def test_prefetch_dimension_contracted_function(self): grid = Grid(shape=(32, 16, 16)) _, _, z = grid.dimensions r = Function(name='r_z', grid=grid, shape=(16,), dimensions=(z,)) kwargs = {'grid': grid, 'space_order': 4, 'time_order': 1, 'save': Buffer(1)} p_x = TimeFunction(name='p_x', **kwargs) p_z = TimeFunction(name='p_z', **kwargs) v = TimeFunction(name='v', **kwargs) v1 = v.func(name='v1') equation = Eq(v.forward, r*p_z.dz + p_x.dx) op0 = Operator(equation, name='contracted_ref', opt=opt0()) op1 = Operator(equation, name='contracted_prefetch', opt=opt4({'par-tile': (16, 8, 8), 'gpu-prefetch': 2})) kernels = [v.root for k, v in op1._func_table.items() if k.startswith('kernel')] staged = [i for kernel in kernels for i in FindNodes(Expression).visit(kernel) if i.write._mem_shared and r in i.reads] assert not staged r.data[:] = 1.0 for f in (p_x, p_z): init_chessboard(f) op0(time_M=2) op1(time_M=2, v=v1) assert np.allclose(v.data, v1.data) def test_par_tile_unfolding(self): grid = Grid(shape=(40, 40, 40)) f = Function(name='f', grid=grid, space_order=8) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=8) eqns = [Eq(g, sin(f)), Eq(u.forward, u.laplace + cos(f).dx)] op = Operator(eqns, opt=opt4({'par-tile': ((32, 8), (64, 4), (32, 4))})) op.cfunction kernels = FindNodes(KernelLaunch).visit(op) assert len(kernels) == 3 assert len({i.block for i in kernels}) == 3 assert [len(i.grid.shape) for i in kernels] == [3, 3, 2] def test_short_par_tile_keeps_outer_dim_on_device(self): """ Pointwise kernels retain 3D blocking with a short `par-tile`. A stencil may iterate the outermost Dimension sequentially, provided the full loop stays inside the GPU kernel. """ grid = Grid(shape=(40, 40, 40)) f = Function(name='f', grid=grid, space_order=8) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=8) eqns = [Eq(g, sin(f)), Eq(u.forward, u.laplace + cos(f).dx)] # The second entry is 2D, while the nest it lands on is 3D op = Operator(eqns, opt=opt4({'par-tile': ((32, 8, 8), (16, 4)), 'gpu-shm': 'min-reuse'})) # The two pointwise kernels unfold their tiles; the stencil keeps a 2D tile launches = FindNodes(KernelLaunch).visit(op) assert [len(i.grid.shape) for i in launches] == [3, 3, 2] x = grid.dimensions[0] kernel = op._func_table['kernel2'].root xloop, = [i for i in FindNodes(Iteration).visit(kernel) if i.dim.root is x] assert xloop.limits == (x.symbolic_min, x.symbolic_max, 1) # No space Dimension is left iterated on the host for i in FindNodes(Iteration).visit(op): assert i.dim.root is grid.time_dim def test_kernelopt_rotation_without_substitution(self): """ An Action that cannot rewrite its Node is not rescheduled for ever. The y-interpolation here sits at a single `x`, so the x-rotation `search_rotation` proposes for it has nothing to substitute and only emits its queue. Without a check the same Action comes back on every pass, growing the Kernel by three expressions a time, and the build never returns. """ grid = Grid(shape=(151, 151)) x, y = grid.dimensions mu = Function(name='mu', grid=grid, space_order=8, staggered=NODE) tau = TimeFunction(name='tau', grid=grid, space_order=8, staggered=(x, y), time_order=1) vx = TimeFunction(name='vx', grid=grid, space_order=8, staggered=x, time_order=1) vy = TimeFunction(name='vy', grid=grid, space_order=8, staggered=y, time_order=1) acc = Function(name='acc', grid=grid, space_order=8, staggered=NODE) mask = INT(Gt(mu, 1e-12)) compliance = tau*SafeInv(mu, mu)*mask/2 + tau*(1 - mask) eqns = [Eq(vx.forward, vx + tau.dx), Eq(acc, acc + compliance*(vx.dy + vy.dx), interp_mode='symmetric')] op = Operator(eqns, subs=grid.spacing_map, opt=opt4({'par-tile': (16, 8, 16)})) TestTemplateKernels._assert_jit(self, op) @pytest.mark.parametrize('space_order', [4, 16]) @pytest.mark.parametrize('independent', [False, True]) @pytest.mark.parametrize('prefetch', [False, 2, 4]) def test_product_queues(self, space_order, independent, prefetch, host_kwargs, extra=None): """Rotate products, retaining an operand queue only for independent reuse.""" grid = Grid(shape=(21, 19, 23), extent=(20., 18., 22.)) a, b, c, d, e = [Function(name=i, grid=grid, space_order=space_order) for i in 'abcde'] u = Function(name='u', grid=grid) v = u.func(name='v') # MFE: the shared `c` used to make the second product fall back to # three operand queues after the first product was placed in a queue expr = (a*b*c).dx + (c*d*e).dx if independent: expr += a.dx eq = Eq(u, expr) op0 = Operator(eq, name='ProductReference', **host_kwargs) op1 = Operator(eq, name='ProductQueues', opt=opt4({ 'cire-mode': '1pass', 'par-tile': (32, 4), 'gpu-prefetch': prefetch, **(extra or {}) })) queues = [f for f in FindSymbols().visit(op1._func_table['kernel0'].root) if isinstance(f, TBQueue)] assert len(queues) == (3 if independent else 2) assert all(q.size == space_order + 1 for q in queues) rng = np.random.default_rng(42) for f in (a, b, c, d, e): f.data_with_halo[:] = rng.uniform(.5, 1.5, f.data_with_halo.shape) op0.apply() op1.apply(u=v) assert np.allclose(u.data, v.data, rtol=1.e-5, atol=1.e-6) @pytest.mark.parametrize('prefetch', [False, 4]) @pytest.mark.parametrize('shifted', [False, True]) def test_reuse_operand_queues(self, prefetch, shifted, host_kwargs): """Preserve operand reuse across products and shifted accesses.""" grid = Grid(shape=(21, 19, 23), extent=(20., 18., 22.)) x = grid.dimensions[0] a, b, c, d = [Function(name=i, grid=grid, space_order=8) for i in 'abcd'] u = Function(name='u', grid=grid) v = u.func(name='v') if shifted: expr = ((a*b).dx(fd_order=2) + a.shift(x, x.spacing).dx + b.shift(x, x.spacing).dx) nqueues = 2 else: xp, xm = x + x.spacing/2, x - x.spacing/2 expr = ((a*b).dx(x0=xp) + (a*c).dx(x0=xm) + (a*d).dx + (b*c).dx(x0=xp) + (b*d).dx(x0=xm) + (c*d).dx) nqueues = 4 eq = Eq(u, expr) op0 = Operator(eq, name='ReuseReference', **host_kwargs) op1 = Operator(eq, name='ReuseQueues', opt=opt4({ 'cire-mode': '1pass', 'par-tile': (32, 4), 'gpu-prefetch': prefetch })) queues = [f for f in FindSymbols().visit(op1._func_table['kernel0'].root) if isinstance(f, TBQueue)] assert len(queues) == nqueues rng = np.random.default_rng(42) for f in (a, b, c, d): f.data_with_halo[:] = rng.uniform(.5, 1.5, f.data_with_halo.shape) op0.apply() op1.apply(u=v) assert np.allclose(u.data, v.data, rtol=1.e-5, atol=1.e-6) class Test2PassAdvanced: def _check(self, op, exp_kernels, exp_queues): assert len(exp_queues) == exp_kernels assert len(FindNodes(KernelLaunch).visit(op)) == exp_kernels for i, expected in enumerate(exp_queues): symbols = FindSymbols().visit(op._func_table[f'kernel{i}'].root) assert len([i for i in symbols if isinstance(i, TBQueue)]) == expected @switchconfig(ignore_unknowns=True) @pytest.mark.parametrize('exprs', [ ('u.dx.dy + u.dy.dz',), ('u.dx.dz + u.dy.dx + u.dy.dz', 'v.dx.dz + v.dy.dx + v.dy.dz'), ('u.dx.dy + u.dy.dz', 'v.dx.dy + v.dy.dz'), ]) @pytest.mark.parametrize('subdomain', [None, 'interior']) def test_mixed_derivs(self, exprs, subdomain): grid = Grid(shape=(100, 37, 50)) subdomain = grid.subdomains.get(subdomain) so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) # noqa u = TimeFunction(name='u', grid=grid, space_order=so) v = u.func(name='v') eqns = [] for expr, field in zip(exprs, [u, v]): eqns.append(Eq(field.forward, eval(expr), subdomain=subdomain)) op0 = Operator(eqns, opt=opt4()) op1 = Operator(eqns, opt=opt4({'cire-mode': 'x|y'})) # Check generated code -- expected only *two* kernels! assert len(FindNodes(KernelLaunch).visit(op1)) == 2 u1 = u.func(name='u1') v1 = u1.func(name='v1') for i in (u, v, u1, v1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, u=u1, v=v1) assert np.allclose(u.data, u1.data, rtol=1e-4) assert np.allclose(v.data, v1.data, rtol=1e-4) def test_flip_noop(self): grid = Grid(shape=(40, 40, 40)) u = TimeFunction(name='u', grid=grid, space_order=4) v = u.func(name='v') eqns = [Eq(u.forward, u.dx.dy + u.dz.dy), Eq(v.forward, v.dx.dy + v.dz.dy)] op0 = Operator(eqns, opt=opt4()) op1 = Operator(eqns, opt=opt4({'cire-mode': 'x|z', 'gpu-flip-derivs': True})) # Check generated code self._check(op1, 2, (2, 0)) u1 = u.func(name='u1') v1 = u1.func(name='v1') for i in (u, v, u1, v1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, u=u1, v=v1) assert np.allclose(u.data, u1.data, rtol=1e-5) assert np.allclose(v.data, v1.data, rtol=1e-5) def test_flip_if_hoisting_and_sparing_registers_v0(self): grid = Grid(shape=(40, 40, 40)) f = Function(name='f', grid=grid) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=4) eqn = Eq(u.forward, (g*(u.dy.dx + f*u.dx) + g*(u.dy.dz + f*u.dy) + g*(u.dx.dz + f*u.dz))) op = Operator(eqn, opt=opt4({'cire-mode': 'x|y', 'gpu-flip-derivs': True})) op.cfunction self._check(op, 2, (1, 0)) def test_flip_if_hoisting_and_sparing_registers_v1(self): grid = Grid(shape=(40, 40, 40)) f = Function(name='f', grid=grid) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=4) eqn = Eq(u.forward, (g.dx*(u.dy.dx + f*u.dx) + g.dy*(u.dy.dz + f*u.dy) + g.dz*(u.dx.dz + f*u.dz))) op = Operator(eqn, opt=opt4({'cire-mode': ('automatic', 'x|y'), 'gpu-flip-derivs': True})) op.cfunction self._check(op, 3, (1, 1, 0)) def test_flip_if_hoisting_and_sparing_registers_v2(self): grid = Grid(shape=(40, 40, 40)) f = Function(name='f', grid=grid) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=4) src = SparseTimeFunction( name='src', npoint=1, coordinates=[(.5, .5, .5)], nt=5, grid=grid ) src_term = src.inject(field=u.forward, expr=src) eqns = [src_term, Eq(u.forward, (g.dx*(u.dy.dx + f*u.dx) + g.dy*(u.dy.dz + f*u.dy) + g.dz*(u.dx.dz + f*u.dz)))] op = Operator(eqns, opt=opt4({'cire-mode': ('automatic', 'automatic', 'x|y'), 'gpu-flip-derivs': True})) op.cfunction self._check(op, 4, (1, 0, 1, 0)) @skipif('nocudalike') @pytest.mark.parametrize('staggered', [False, True]) @pytest.mark.parametrize('cire_mode,selected', [ ('x:stag', (1, 5)), ('x:center', (0, 6)), ('x:interp', (2, 4)), ('x', (0, 1, 2, 3, 4, 5, 6, 7)), ('x2:center', (8,)), ('x2:stag', (9,)) ]) def test_cire_kinds(self, cire_mode, selected, staggered): grid = Grid(shape=(31, 17, 25)) x, _, _ = grid.dimensions fields = [TimeFunction(name=f'f{i}', grid=grid, space_order=16, staggered=x if staggered else NODE) for i in range(10)] out = TimeFunction(name='out', grid=grid, space_order=16) out1 = out.func(name='out1') # Evaluation locations are relative to the input lattice, including x+h/2 x0 = fields[0].indices[1] exprs = [fields[0].dx(x0={x: x0}), fields[1].dx(x0={x: x0 + x.spacing/2}), fields[2].diff(x, deriv_order=0, fd_order=16, x0={x: x0 + x.spacing/2}), fields[3].dx(x0={x: x0 + x.spacing/4}), fields[4].diff(x, deriv_order=0, fd_order=8, x0={x: x0 + x.spacing/2}), # A two-point derivative has a half-cell stencil centre even # when requested on-grid; custom weights can override x0 too fields[5].dx(fd_order=2, x0={x: x0}), fields[6].dx(weights=(-1, 0, 1), x0={x: x0 + x.spacing/4}), fields[7].dx(weights=(-2, 1, 1), x0={x: x0}), fields[8].dx2(fd_order=4, x0={x: x0}), fields[9].dx2(fd_order=4, x0={x: x0 + x.spacing/2})] eqn = Eq(out.forward, out + sum(exprs)) op0 = Operator(eqn, name='op0', opt=opt0({'cire-mode': '1pass'})) op1 = Operator(eqn, name='op1', opt=opt4({'cire-mode': cire_mode, 'gpu-flip-derivs': True})) producer, _ = FindNodes(KernelLaunch).visit(op1) assert tuple(i for i, f in enumerate(fields) if f.dmap in producer.arguments) == selected for f in fields: init_chessboard(f) op0.apply(time_M=2) op1.apply(time_M=2, out=out1) assert np.allclose(out.data, out1.data, rtol=1.e-4, atol=1.e-5) @skipif('nocudalike') @pytest.mark.parametrize('flip', [False, True]) @pytest.mark.parametrize('cire_mode,expected', [ ('x|y|z:stag,y|z:center,x|y|z:interp', 4), ('x|y|z:interp,y|z:center,x|y|z:stag', 2), ('x|y|z:stag,x|y|z:stag,y|z:center,z:interp', 4), ('x|y|z:stag,x|y|z:stag,x|y|z:stag,y|z:center,z:interp', 4), ('1pass', 1) ]) def test_cire_passes(self, cire_mode, expected, flip): grid = Grid(shape=(31, 17, 25)) x, y, z = grid.dimensions u = TimeFunction(name='u', grid=grid, space_order=16) v = TimeFunction(name='v', grid=grid, space_order=16) v1 = v.func(name='v1') expr = u.dx(fd_order=8, x0={x: x + x.spacing/2}).dy(fd_order=8) expr = expr.diff(z, deriv_order=0, fd_order=16, x0={z: z + z.spacing/2}) eqn = Eq(v.forward, v + expr) op0 = Operator(eqn, name='op0', opt=opt0({'cire-mode': '1pass'})) op1 = Operator(eqn, name='op1', opt=opt4({'cire-mode': cire_mode, 'gpu-flip-derivs': flip})) assert len(FindNodes(KernelLaunch).visit(op1)) == expected init_chessboard(u) u.data_with_halo[:] *= 0.01 op0.apply(time_M=2) op1.apply(time_M=2, v=v1) assert np.allclose(v.data, v1.data, rtol=1.e-4, atol=1.e-5) @skipif('nocudalike') @pytest.mark.parametrize('flip', [False, True]) @pytest.mark.parametrize('cire_mode,expected,exp_shmas', [ ('xy:interp', (2, 2), (0, 0)), ('x|y:interp', (2, 2), (1, 1)), ('x|y:interp,x|y:interp', (3, 3), (0, 0)), ('x:interp,y:interp', (3, 3), (0, 0)), ('y:interp,x:interp', (2, 3), (1, 0)) ]) def test_cire_interpolation(self, cire_mode, expected, exp_shmas, flip): grid = Grid(shape=(31, 17, 25)) x, y, _ = grid.dimensions u = TimeFunction(name='u', grid=grid, space_order=16) v = TimeFunction(name='v', grid=grid, space_order=16) v1 = v.func(name='v1') expr = u.diff(x, deriv_order=0, fd_order=16, x0={x: x + x.spacing/2}) expr = expr.diff(y, deriv_order=0, fd_order=16, x0={y: y + y.spacing/2}) eqn = Eq(v.forward, v + expr) op0 = Operator(eqn, name='op0', opt=opt0({'cire-mode': '1pass'})) op1 = Operator(eqn, name='op1', opt=opt4({'cire-mode': cire_mode, 'par-tile': (16, 16), 'gpu-flip-derivs': flip})) # Flip can expose y first; otherwise that pass finds no inner y stencil launches = FindNodes(KernelLaunch).visit(op1) assert len(launches) == expected[flip] # A remaining y interpolation needs shared memory in the final consumer kernel = op1._func_table[launches[-1].name].root shmas = [i for i in FindSymbols().visit(kernel) if i._mem_shared] assert len(shmas) == exp_shmas[flip] init_chessboard(u) op0.apply(time_M=2) op1.apply(time_M=2, v=v1) assert np.allclose(v.data, v1.data, rtol=1.e-4, atol=1.e-5) @skipif('nocudalike') @pytest.mark.parametrize('trace', [False, True]) @pytest.mark.parametrize('schedule,expected', [(None, 6), (0, 9)]) def test_cire_strain(self, schedule, expected, trace): grid = Grid(shape=(31, 17, 25)) v = VectorTimeFunction(name='v', grid=grid, space_order=8, interp_order=8) tau = TensorTimeFunction(name='tau', grid=grid, space_order=8, interp_order=8) tau1 = TensorTimeFunction(name='tau1', grid=grid, space_order=8, interp_order=8) mu = Constant(name='mu', value=0.01) g = grad(v) strain = (g + g.transpose(inner=False))/2 stress = mu*strain if trace: lam = Constant(name='lam', value=0.03) stress += lam*diag(strain.trace()) op0 = Operator(Eq(tau.forward, tau + stress), name='op0', opt=opt0({'cire-mode': '1pass'})) op1 = Operator(Eq(tau1.forward, tau1 + stress), name='op1', opt=opt4({'cire-mode': 'x|y|z:stag', 'cire-schedule': schedule})) arrays = [i for i in FindSymbols().visit(op1.body) if isinstance(i, TempArray)] assert len(arrays) == expected for f in v: init_chessboard(f) op0.apply(time_M=2) op1.apply(time_M=2) for f, f1 in zip(tau.values(), tau1.values()): assert np.allclose(f.data, f1.data, rtol=1.e-4, atol=1.e-5) class TestWriteAheadStores: @pytest.mark.parametrize('options', [ {}, {'cire-mode': 'x|y'}, ]) def test_basic(self, options): grid = Grid(shape=(31, 17, 25)) x, _, _ = grid.dimensions sources = [Function(name=f'source{i}', grid=grid, space_order=4) for i in range(2)] transverse = [Function(name=f'transverse{i}', grid=grid, space_order=4) for i in range(2)] longitudinal = [Function(name=f'longitudinal{i}', grid=grid, space_order=4) for i in range(2)] eqns = [] for source, ft, fl in zip(sources, transverse, longitudinal): eqns.extend([Eq(ft, source.dy2), Eq(fl, source.dx2 + source.dy.dx + source.dy.dz)]) options = {'par-tile': (10, 7, 5), **options} op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt4(options)) kernels = [v.root for k, v in op1._func_table.items() if k.startswith('kernel')] for f in transverse: writes = [i.expr.lhs for kernel in kernels for i in FindNodes(Expression).visit(kernel) if i.write.function is f] assert len(writes) == 2 assert any(str(i.indices[x]) == 'x + 6' for i in writes) assert any(any(isinstance(s, QueueInitDimension) for s in i.free_symbols) for i in writes) outputs = transverse + longitudinal references = [f.func(name=f'{f.name}_ref') for f in outputs] for source in sources: init_chessboard(source) op0.apply() op1.apply(**{f.name: ref for f, ref in zip(outputs, references)}) for f, ref in zip(outputs, references): assert np.allclose(f.data, ref.data) def test_strain_like(self): grid = Grid(shape=(31, 17, 25)) x, _, _ = grid.dimensions so = (16, 8, 8) v = VectorFunction(name='v', grid=grid, space_order=so, staggered=(NODE, NODE, NODE)) v_x, v_y, v_z = v.values() e = TensorFunction(name='e', grid=grid, space_order=so, staggered=NODE) e_xx, e_xy, e_xz, e_yy, e_yz, e_zz = e.values() inputs = (v_x, v_y, v_z) yz_only = (e_yy, e_yz, e_zz) x_dependent = (e_xx, e_xy, e_xz) # Compact high-order analogue of the six-field strain kernel eqns = [] for source, fyz, fx in zip(inputs, yz_only, x_dependent): eqns.extend([Eq(fyz, source.dy2), Eq(fx, source.dx2 + source.dy.dx + source.dy.dz)]) op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt4({'cire-mode': 'x|y', 'par-tile': (16, 8, 16)})) kernels = [v.root for k, v in op1._func_table.items() if k.startswith('kernel')] for f in yz_only: writes = [i.expr.lhs for kernel in kernels for i in FindNodes(Expression).visit(kernel) if i.write.function is f] assert len(writes) == 2 assert any(str(i.indices[x]) == 'x + 16' for i in writes) assert any(any(isinstance(s, QueueInitDimension) for s in i.free_symbols) for i in writes) outputs = x_dependent + yz_only references = [f.func(name=f'{f.name}_ref') for f in outputs] for f in inputs: init_chessboard(f) op0.apply() op1.apply(**{f.name: ref for f, ref in zip(outputs, references)}) for f, ref in zip(outputs, references): assert np.allclose(f.data, ref.data) def test_guarded_interpolations(self): grid = Grid(shape=(31, 17, 25)) so = (16, 8, 8) n = VectorFunction(name='n', grid=grid, space_order=so) e = TensorFunction(name='e', grid=grid, space_order=so, interp_order=16) v = VectorTimeFunction(name='v', grid=grid, space_order=so, time_order=1) tau = TensorTimeFunction(name='tau', grid=grid, space_order=so, time_order=1) g = grad(v.forward) e_expr = (g + g.transpose(inner=False))/2 w = e*n stress = n*w.T + w*n.T eqns0 = [Eq(lhs, current + rhs) for lhs, current, rhs in zip(tau.forward.diagonal(1), tau.diagonal(1), stress.diagonal(1))] eqns = [Eq(e, e_expr), *eqns0] op0 = Operator(eqns, opt=opt0({'errctl': 'basic'})) op1 = Operator(eqns, opt=opt4({ 'errctl': 'basic', # Leave the z derivatives in the interpolation consumer 'cire-mode': ('1pass', 'x|y'), 'par-tile': (16, 8, 16), })) kernels = [v.root for k, v in op1._func_table.items() if k.startswith('kernel')] queues = {i for kernel in kernels for i in FindSymbols().visit(kernel) if isinstance(i, TBQueue)} assert len(queues) == 9 outputs = [*e.values(), *tau.diagonal(1)] references = [f.func(name=f'{f.name}_ref') for f in outputs] for f in (*n.values(), *v.values(), *outputs, *references): init_chessboard(f) op0.apply(time_M=2) op1.apply(time_M=2, **{f.name: ref for f, ref in zip(outputs, references)}) for f, ref in zip(outputs, references): assert np.allclose(f.data, ref.data) class TestGlobalLoadsHoisting: @pytest.mark.parametrize('ndim, expected_kernels', [(2, 2), (3, 1)]) def test_forked_nest(self, ndim, expected_kernels): grid = Grid(shape=(4,)*ndim) *outer, d = grid.dimensions dl = SubDimension.left(name='dl', parent=d, thickness=1, separated=True) dr = SubDimension.right(name='dr', parent=d, thickness=1, separated=True) p = Function(name='p', grid=grid, space_order=0) # The 3D case retains two shared parallel loops and keeps the fork # intact. Hoisted loads must stay in the same loop as their consumers left, right = p[(*outer, dl)], p[(*outer, dr)] eqns = [Eq(left, left + 1), Eq(right, right + 1)] op = Operator(eqns, name='hoist_subdomains', opt=('advanced', { 'gpu-opt': True, 'gpu-hoist': 2, 'fission': False })) launches = FindNodes(KernelLaunch).visit(op) assert len(launches) == expected_kernels assert all(len(i.block.cargs) == 2 for i in launches) p.data[:] = np.arange(4**ndim, dtype=np.float32).reshape(grid.shape) expected = p.data.copy() expected[..., (0, -1)] += 1 op.apply() assert np.array_equal(p.data, expected) @pytest.mark.parametrize('expr,extra,expected', [ ('u.dy.dz + u + u.backward', {'gpu-hoist': True}, [('float r0 = uL0(t1, x + 4, y + 4, z + 2)', -4)]), ('u.dy.dz + u + u.backward', {'gpu-hoist': 2}, [('r0 = uL0(t1, x + 4, y + 4, z + 2)', 10)]), ]) def test_basic(self, expr, extra, expected): grid = Grid(shape=(10, 10, 10)) nt = 5 u = TimeFunction(name='u', grid=grid, space_order=4, time_order=2) eqn = Eq(u.forward, eval(expr) + .01) op0 = Operator(eqn, opt=opt4()) op1 = Operator(eqn, opt=opt4(extra)) # Check generated code kernel = op1._func_table['kernel0'].root exprs = FindNodes(Expression).visit(kernel) for e, i in expected: assert f'{e};' == str(exprs[i]) u1 = u.func(name='u1') init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) def test_candidate_identical_to_write(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4) f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) eqns = [Eq(f, f + g + u + 2.), Eq(g, u + 1)] op = Operator(eqns, opt=opt4({'par-tile': (4, 4), 'gpu-hoist': 1})) op.apply(time_M=9) assert np.all(g.data == 1) assert np.all(f.data == 29) def test_iteration_indep_raw_dependence(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4) f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) eqns = [Eq(f, u + 2.), Eq(g, g + f + u)] op = Operator(eqns, opt=opt4({'par-tile': (4, 4), 'gpu-hoist': 1})) # Check generated code kernel = op._func_table['kernel0'].root exprs = FindNodes(Expression).visit(kernel) assert len(exprs) == 6 assert exprs[2].write.name == 'is_writer_glb' assert str(exprs[3]) == 'float r0 = gL0(x + 4, y + 4, z + 4);' assert 'r1' not in str(kernel) op.apply(time_M=9) assert np.all(f.data == 2) assert np.all(g.data == 20) def test_buffer_1(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4, time_order=1, save=Buffer(1)) v = TimeFunction(name='v', grid=grid, space_order=4, time_order=1, save=Buffer(1)) eqn = Eq(u.forward, v.laplace + u + 2.) op = Operator(eqn, opt=opt4({'par-tile': (4, 4), 'gpu-hoist': 1})) op.cfunction # Check generated code kernel = op._func_table['kernel0'].root exprs = FindNodes(Expression).visit(kernel) assert 'float r0 = uL0(0, x + 4, y + 4, z + 4)' in str(exprs[-8]) class TestKernelOptVectTypes: _subdomain_specs = [ ('left', 'left', 'left'), ('right', 'right', 'right'), ('middle', 'middle', 'middle'), ('left', 'right', 'middle') ] @pytest.mark.parametrize('spec', _subdomain_specs) @pytest.mark.parametrize('override', [False, True]) def test_no_bundle_subdomains(self, spec, override): grid = Grid(shape=(10, 12, 14)) sd = SpeccedDomain('boundary', spec, 2, grid=grid) h = Function(name='h', grid=grid, space_order=2) f = Function(name='f', grid=sd, space_order=2) g = Function(name='g', grid=sd, space_order=2) eqns = [Eq(h, h + 1, subdomain=grid.interior), Eq(f, f + g + 1)] op = Operator(eqns, name='no_bundle_subdomains', opt=opt5()) assert not any(f.is_Bundle for f in FindSymbols().visit(op)) # Also exercise different allocation sizes if override: grid = Grid(shape=(11, 13, 17)) sd = SpeccedDomain('boundary', spec, 2, grid=grid) h = Function(name='h1', grid=grid, space_order=2) f = Function(name='f1', grid=sd, space_order=2) g = Function(name='g1', grid=sd, space_order=2) f_ref = f.func(name='f_ref') for i in (f, f_ref, g): init_chessboard(i) g.data_with_halo[:] *= 2 f_ref.data[:] += g.data + 1 expected_g = np.array(g.data_with_halo) expected_h = np.zeros_like(h.data) expected_h[1:-1, 1:-1, 1:-1] = 1 op.apply(h=h, f=f, g=g) np.testing.assert_array_equal(f.data_with_halo, f_ref.data_with_halo) np.testing.assert_array_equal(g.data_with_halo, expected_g) np.testing.assert_array_equal(h.data, expected_h) @pytest.mark.parametrize('spec', _subdomain_specs) @pytest.mark.parametrize('override', [False, True]) def test_no_bundle_subdomains_time(self, spec, override): grid = Grid(shape=(10, 12, 14)) sd = SpeccedDomain('boundary', spec, 2, grid=grid) h = Function(name='h', grid=grid, space_order=2) u = TimeFunction(name='u', grid=sd, space_order=2) v = TimeFunction(name='v', grid=sd, space_order=2) eqns = [Eq(h, h + 1, subdomain=grid.interior), Eq(u.forward, u + v + 1)] op = Operator(eqns, name='no_bundle_subdomains_time', opt=opt5()) assert not any(f.is_Bundle for f in FindSymbols().visit(op)) if override: grid = Grid(shape=(11, 13, 17)) sd = SpeccedDomain('boundary', spec, 2, grid=grid) h = Function(name='h1', grid=grid, space_order=2) u = TimeFunction(name='u1', grid=sd, space_order=2) v = TimeFunction(name='v1', grid=sd, space_order=2) u_ref = u.func(name='u_ref') for i in (u, u_ref, v): init_chessboard(i) v.data_with_halo[:] *= 2 u_ref.data[1] = u.data[0] + v.data[0] + 1 expected_v = np.array(v.data_with_halo) expected_h = np.zeros_like(h.data) expected_h[1:-1, 1:-1, 1:-1] = 1 op.apply(h=h, u=u, v=v, time_M=0) np.testing.assert_array_equal(u.data_with_halo, u_ref.data_with_halo) np.testing.assert_array_equal(v.data_with_halo, expected_v) np.testing.assert_array_equal(h.data, expected_h) def test_bundle_full_grid_in_subdomain(self): grid = Grid(shape=(10, 12, 14)) f = Function(name='f', grid=grid, space_order=2) g = Function(name='g', grid=grid, space_order=2) f_ref = f.func(name='f_ref') op = Operator(Eq(f, f + g + 1, subdomain=grid.interior), name='bundle_full_grid_in_subdomain', opt=opt5()) bundles = [f for f in FindSymbols().visit(op) if f.is_Bundle] assert len(bundles) == 1 assert bundles[0].components == (f, g) for i in (f, f_ref, g): init_chessboard(i) g.data_with_halo[:] *= 2 f_ref.data[1:-1, 1:-1, 1:-1] += g.data[1:-1, 1:-1, 1:-1] + 1 expected_g = np.array(g.data_with_halo) op.apply() np.testing.assert_array_equal(f.data_with_halo, f_ref.data_with_halo) np.testing.assert_array_equal(g.data_with_halo, expected_g) def test_bundle_stencil_gap(self): grid = Grid(shape=(8, 8, 24)) z = grid.dimensions[-1] zl = SubDimension.left('zl', z, 8, separated=True) zr = SubDimension.right('zr', z, 8, separated=True) f = Function(name='f', grid=grid, space_order=8) g = Function(name='g', grid=grid, space_order=8) g.data[:] = 2 eq = Eq(f, f + g + 1) op = Operator([eq.subs(z, zl), eq.subs(z, zr)], name='bundle_stencil_gap', opt=opt5()) bundles = [i for i in FindSymbols().visit(op) if i.is_Bundle] assert len(bundles) == 1 assert bundles[0].components == (f, g) op.apply() expected = np.zeros(grid.shape) expected[:, :, :8] = expected[:, :, -8:] = 3 np.testing.assert_array_equal(f.data, expected) assert np.all(g.data == 2) runtime_grid = Grid(shape=(8, 8, 23)) fields = {i.name: Function(name=f'{i.name}_runtime', grid=runtime_grid, space_order=8) for i in (f, g)} with pytest.raises(InvalidArgument, match='gap of at least 8'): op.arguments(**fields) def test_vt_in_isolation_v0(self): grid = Grid(shape=(11,)) f = Function(name='f', grid=grid) g = Function(name='g', grid=grid) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) init_chessboard(u) init_chessboard(u1) init_chessboard(v) eqn = Eq(u.forward, u.dx + v.dx + f + g + 1) op0 = Operator(eqn, opt=opt0()) op1 = Operator(eqn, opt=opt5({'par-tile': (3,)})) # Check generated code -- vector types are being used! assert 'float2' in str(op1) assert 'uvL0(t1, x + 4).x =' in str(op1) assert 'fg[x + 1].x + fg[x + 1].y' in str(op1) op0.apply(time_M=10) op1.apply(time_M=10, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-6) def test_vt_in_isolation_v1(self): grid = Grid(shape=(8, 8), dtype=np.int32) x, y = grid.dimensions t = grid.stepping_dim # noqa u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) v1 = TimeFunction(name='v', grid=grid, space_order=4) eqns = [ Eq(u.forward, v.dx2 + u.dy2 + v + 1), Eq(v.forward, u.dx2 + v.dy2 + v + 1), ] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt5()) # Check generated code -- vector types are being used! assert 'int2' in str(op1) assert 'uvL0(t1, x + 4, y + 4) = make_int2(' in str(op1) op0.apply(time_M=4) op1.apply(time_M=4, u=u1, v=v1) assert np.all(u.data == u1.data) assert np.all(v.data == v1.data) @pytest.mark.parametrize('shape', [(21, 11, 89)]) def test_vt_in_isolation_v2(self, shape): grid = Grid(shape=shape) u = TimeFunction(name='u', grid=grid, space_order=4) f = Function(name='f', grid=grid) g = Function(name='g', grid=grid) h = Function(name='h', grid=grid) eq = Eq(f, f + g + h + u.dx2) op = Operator(eq, opt=opt5()) # Check generated code assert 'fghL0(x + 1, y + 1, z + 1).x =' in str(op) g.data[:] = 1.2 h.data[:] = 0.3 op.apply(time_M=0) assert np.all(f.data == 1.5) def test_vt_in_isolation_v3(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) f = Function(name='f', grid=grid) g = Function(name='g', grid=grid) h = Function(name='h', grid=grid) eqns = [Eq(u.forward, f + g + u.dx2), Eq(v.forward, v*h + v.dx2)] op = Operator(eqns, opt=opt5()) # Check generated code assert 'fghL0(x + 1, y + 1, z + 1).x' in str(op) assert 'uvL0(t, x + 4, y + 4, z + 4) =' in str(op) assert len(op._func_table) == 11 assert 'unbundle1' not in op._func_table # Do not unbundle read-only assert 'free1' not in op._func_table @pytest.mark.parametrize('shape', [(21, 11, 89)]) def test_vt_in_isolation_complex0(self, shape, host_kwargs): nt = 4 grid = Grid(shape=shape) time_dim = grid.time_dim factor = Constant(name='factor', value=2, dtype=np.int32) time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) v1 = TimeFunction(name='v', grid=grid, space_order=4) usave = TimeFunction(name='usave', grid=grid, time_order=0, save=(int(nt//factor.data)), time_dim=time_sub, layers=NoLayers) usave1 = usave._rebuild() sf = SparseTimeFunction(name='sf', grid=grid, nt=nt, npoint=1, coordinates=[(.5, .5, .5)]) eqns = [ Eq(u.forward, u.dx.dy + u.dy.dz + u.dx.dz + v + .01), Eq(v.forward, v.dx.dy + v.dy.dz + v.dx.dz + u + .01), sf.inject(field=u, expr=sf + 1), Eq(usave, u), ] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt5()) # Check generated code -- vector types are being used! assert 'float2' in str(op1) assert 'uvL0(t1, x + 4, y + 4, z + 4) = make_float2' in str(op1) op0.apply(time_M=nt-1) op1.apply(time_M=nt-1, u=u1, v=v1, usave=usave1) # Numerics diverge quickly, early loss of accuracy, hence the low # tolerances. Look at the other tests above for tighther tolerances assert np.allclose(u.data, u1.data, rtol=1.e-2) assert np.allclose(v.data, v1.data, rtol=1.e-2) assert np.allclose(usave.data, usave1.data, rtol=1.e-6) def test_complete0(self, host_kwargs): grid = Grid(shape=(18, 34, 21)) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) init_chessboard(u) init_chessboard(u1) init_chessboard(v) eqns = [Eq(u.forward, u.dz2 + u.dy.dz + 1.), Eq(v.forward, v.dz2 + v.dy.dz + 1.)] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt6({'par-tile': (15, 5, 2)})) # Check generated code -- vector types, rotations, and shared memory are # being used! assert all(i in str(op1) for i in ['float2', 'uvL0']) op0.apply(time_M=5) op1.apply(time_M=5, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) def test_complete1(self, host_kwargs): grid = Grid(shape=(21, 53, 22)) f = Function(name='f', grid=grid) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) init_chessboard(u) init_chessboard(u1) init_chessboard(v) eqns = [ Eq(u.forward, u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz + 1.), Eq(v.forward, v.dx2 + v.dy2 + v.dz2 + v.dy.dx + v.dy.dz + f*v.dx.dz + 1.) ] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt6({'par-tile': (16, 4)})) # Check generated code -- vector types, rotations, and shared memory are # being used! assert all(i in str(op1) for i in ['float2', 'uvL0']) efunc = op1._func_table['kernel0'].root # Make sure there's the right amount of fusion! assert len(FindNodes(Conditional).visit(efunc)) == 8 op0.apply(time_M=5) op1.apply(time_M=5, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) def test_issue_400(self): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) h = Function(name='h', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, 4*f*g + 4*h + u.dx + v + 1), Eq(v.forward, 4*g*h + 4*f + v.dx + u + 1)] op = Operator(eqns, opt=opt6()) assert "make_float2" in str(op._func_table['kernel0'].root) section = op._profiler._sections['section0'] assert section.sops == 32 def test_vect_minimal(self): grid = Grid(shape=(18, 34, 21)) f = Function(name='f', grid=grid) g = Function(name='g', grid=grid) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dz2 + u.dx2 + f*g + 1.), Eq(v.forward, v.dz2 + v.dx2 + f + g + 1.)] op = Operator(eqns, opt=opt6({'gpu-vect': 'minimal'})) assert all(i in str(op) for i in ['float2', 'fgL0']) assert 'uvL0' not in str(op) queues = [i for i in FindSymbols().visit(op._func_table['kernel0'].root) if i.name.startswith('queue')] assert len(queues) == 2 @pytest.mark.parallel(mode=[(4, 'diag2')]) @pytest.mark.parametrize('exprs', [ ('u.dx2 + u', 'v.dx2 + v'), ('u.dx2 + u.dy + v.dz', 'v.dx2 + v.dy + u.dz') ]) def test_vt_in_isolation_mpi(self, exprs, mode): grid = Grid(shape=(37, 48, 21)) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) v1 = TimeFunction(name='v', grid=grid, space_order=4) eqns = [] for expr, field in zip(exprs, [u, v]): eqns.append(Eq(field.forward, eval(expr) + 1.)) u.data_with_halo[:] = 1. u1.data_with_halo[:] = 1. v.data_with_halo[:] = 1. v1.data_with_halo[:] = 1. op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt5()) # Check generated code assert 'haloupdate0(uv_vec' in str(op1) op0.apply(time_M=1) op1.apply(time_M=1, u=u1, v=v1) assert np.all(u.data == u1.data) assert np.all(v.data == v1.data) def test_w_mixed_precision(self): grid = Grid(shape=(10, 10, 10)) # Dummy values for mixed precision -- not gonna checke the output value # The aims of the test are that the generated code is correct and that # the Operator can run without errors m, M = 0, 1 p = TimeFunction(name='p', grid=grid, space_order=4, time_order=2) # 2 x float32 f0_32 = Function(name='f0_32', grid=grid, space_order=4) f1_32 = Function(name='f1_32', grid=grid, space_order=4) # 1 x float16 float16 = Float16(m, M) f0_16 = Function(name='f0_16', grid=grid, space_order=4, dtype=float16) # 4 x float8 float8 = Float8(m, M) f0_8 = Function(name='f0_8', grid=grid, space_order=4, dtype=float8) f1_8 = Function(name='f1_8', grid=grid, space_order=4, dtype=float8) f2_8 = Function(name='f2_8', grid=grid, space_order=4, dtype=float8) f3_8 = Function(name='f3_8', grid=grid, space_order=4, dtype=float8) eqn = Eq(p.forward, p.dx*f0_32*f1_32*f0_16*f0_8*f1_8*f2_8*f3_8) op = Operator(eqn, opt=opt6()) # Check generated code assert 'float2 *f0_32f1_32' in str(op) # one vect type for the two float32 assert 'unsigned short *f0_16' in str(op) # no vect types for float16 assert 'unsigned char *f0' in str(op) # no vect types for float8 assert 'f0_8f0_16' not in str(op) op.apply(time_M=1) def test_sin_cos(self): grid = Grid(shape=(10, 10, 10)) f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4) eq = Eq(u.forward, u.dx2 + sin(f) + cos(f) + sin(g) + cos(g)) op = Operator(eq, opt=opt6()) assert '__sincos' in str(op) assert not any(i in str(op) for i in ('sin(', '__sin(', ' cos(', ' __cos(')) op.cfunction def test_sin_cos_cso(self): grid = Grid(shape=(10, 10)) u = TimeFunction(name='u', grid=grid, space_order=2) theta = Function(name='theta', grid=grid, dtype=Float16(0, 1.5)) eq = Eq(u.forward, sin(theta) * u.dx2 + cos(theta) * u.dy2) op = Operator(eq) # First temp should be the cso unpack assert 'float r0 = f160 +' in str(op) # sincos uses r0 assert '__sincosf(r0,&r3,&r4)' in str(op) # Make sure it compiles as expected with the right order of operations op.cfunction @pytest.mark.parametrize('space_orders', [ (8, 8, 8, 8), (16, 16, 16, 8), (16, 16, 16, 16), (32, 16, 16, 16), (32, 32, 32, 32), (16, 32, 32, 16) ]) def test_bundle_w_varying_space_order(self, space_orders): grid = Grid(shape=(10, 10, 10)) so0, so1, so2, so3 = space_orders u = TimeFunction(name='u', grid=grid, space_order=so0) f0 = Function(name='f0', grid=grid, space_order=so1) f1 = Function(name='f1', grid=grid, space_order=so2) f2 = Function(name='f2', grid=grid, space_order=so3) eq = Eq(u.forward, u.dx + f0 + f1 + f2) op = Operator(eq, opt=opt6()) op.cfunction def test_staggered_stencils_v0(self): grid = Grid(shape=(10, 10, 10)) z = grid.dimensions[-1] u = TimeFunction(name='u', grid=grid, space_order=4, staggered=z) v = TimeFunction(name='v', grid=grid, space_order=4, staggered=NODE) eqn = Eq(u.forward, u.dz + v.dz + 1) op = Operator(eqn, opt=opt6()) # Different stencils, no vector types assert not any('float%d' % i in str(op) for i in range(2, 5)) def test_staggered_stencils_v1(self): grid = Grid(shape=(10, 10, 10)) z = grid.dimensions[-1] u0 = TimeFunction(name='u0', grid=grid, space_order=4, staggered=z) u1 = TimeFunction(name='u1', grid=grid, space_order=4, staggered=z) v = TimeFunction(name='v', grid=grid, space_order=4, staggered=NODE) eqn = Eq(u1.forward, u0.dz + u1.dz + v.dz + 1) op = Operator(eqn, opt=opt6()) # Expected float2 for u0 and u1, but not for v assert 'float2 *u0u1' in str(op) assert not any('float%d' % i in str(op) for i in range(3, 5)) def test_staggered_stencils_v2(self): grid = Grid(shape=(10, 10, 10)) z = grid.dimensions[-1] u0 = TimeFunction(name='u0', grid=grid, space_order=4, staggered=z) u1 = TimeFunction(name='u1', grid=grid, space_order=4, staggered=z) u2 = TimeFunction(name='u2', grid=grid, space_order=4, staggered=z) v0 = TimeFunction(name='v0', grid=grid, space_order=4, staggered=NODE) v1 = TimeFunction(name='v1', grid=grid, space_order=4, staggered=NODE) v2 = TimeFunction(name='v2', grid=grid, space_order=4, staggered=NODE) eqn = Eq(u1.forward, (u0.dz + u1.dz + u2.dz + v0.dz + v1.dz + v2.dz)) op = Operator(eqn, opt=opt6()) assert 'u0u1u2' in str(op) assert 'v0v1v2' in str(op) assert not any('float%d' % i in str(op) for i in [2, 4]) def test_staggered_stencils_v3(self): grid = Grid(shape=(10, 10, 10)) z = grid.dimensions[-1] f = Function(name='f', grid=grid, space_order=4, parameter=True) u = TimeFunction(name='u', grid=grid, space_order=4, staggered=z) g = Function(name='g', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dz*f), Eq(v.forward, v.dz*g)] op = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) assert 'float2 *uv' in str(op) assert 'float2 *fg' not in str(op) def test_staggered_stencils_v4(self): grid = Grid(shape=(10, 10, 10)) z = grid.dimensions[-1] f0 = Function(name='f0', grid=grid, space_order=4, parameter=True) f1 = Function(name='f1', grid=grid, space_order=4, parameter=True) u = TimeFunction(name='u', grid=grid, space_order=4, staggered=z) g0 = Function(name='g0', grid=grid, space_order=4) g1 = Function(name='g1', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dz*f0*f1), Eq(v.forward, v.dz*g0*g1)] op = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) assert 'float2 *d_uv' in str(op) assert 'float2 *d_f0f1' in str(op) assert 'float2 *d_g0g1' in str(op) @pytest.mark.parametrize('exprs', [ ('tau[0].dx', 'v.forward[0].dx'), ('tau[0].dx + tau[0].dy', 'v.forward[0].dx'), ('tau[0].dy', 'v.forward[0].dx'), ('tau[0].dx + tau[0].dy', 'v.forward[0].dx + v.forward[0].dy'), ]) @pytest.mark.parallel(mode=[(4, 'diag2')]) def test_staggered_stencils_v5_mpi(self, exprs, mode): grid = Grid(shape=(60, 60)) v = VectorTimeFunction(name='v', grid=grid, space_order=4, time_order=0) tau = TensorTimeFunction(name='tau', grid=grid, space_order=4, time_order=0) eqns = [Eq(v.forward, v - eval(exprs[0]) + .00001), Eq(tau.forward, tau - eval(exprs[1]))] op0 = Operator(eqns, opt=opt4({'par-tile': (10, 10)})) op1 = Operator(eqns, opt=opt6({'par-tile': (10, 10)})) # Check generated code assert all(i in str(op1) for i in ['float2', 'float3']) # NOTE: the IET engine should avoid the generation of redundant # (i.e. semantically identical) callables assert all(i in op1._func_table for i in ['haloupdate0', 'halowait0']) assert len([i for i in op1._func_table.values() if isinstance(i.root, (HaloUpdate, HaloWait))]) == 2 fields = tau.values() + v.values() overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} op0.apply(time_M=2) op1.apply(time_M=2, **overrides) for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) def test_tbarrays_v0(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dx.dy), Eq(v.forward, v.dx.dz)] op = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) assert 'float2 queue2' in str(op) assert 'float queue' not in str(op) def test_two_pass(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=16) v = TimeFunction(name='v', grid=grid, space_order=16) eqns = [Eq(u.forward, u.dy.dz), Eq(v.forward, v.dy.dz)] op = Operator(eqns, opt=opt6({'par-tile': (4, 4), 'cire-mode': '2pass'})) op.cfunction assert 'float2 *r0r1' in str(op) assert 'float2 *d_uv' in str(op) def test_conflicting_kernels_v0(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) w = TimeFunction(name='w', grid=grid, space_order=4) o = TimeFunction(name='o', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dx + v.dx + w.dx + 1.), Eq(v.forward, u.dx + v.dx + w.dx + 2.), Eq(w.forward, u.dx + v.dx + w.dx + 3.), Eq(o.forward, u.forward.dx + v.forward.dx + w.forward.dy)] op = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) assert 'float3 *restrict uvw' in str(op) assert 'make_float3' in str(op) def test_is_transient(self): grid = Grid(shape=(30, 30, 30)) u = TimeFunction(name='u', grid=grid, is_transient=True) v = TimeFunction(name='v', grid=grid, is_transient=True) w = TimeFunction(name='w', grid=grid, save=Buffer(1)) w1 = w.func(name='w1') eqns = [Eq(u.forward, u.dx + v.dx + 1.), Eq(v.forward, u.dx + v.dx + 2.), Eq(w, u.forward + v.forward)] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) op0.apply(time_M=2) op1.apply(time_M=2, w=w1) # Check generated code assert 'float2 *d_uv' in str(op1) assert 'bundle' not in str(op1) assert 'unbundle' not in str(op1) kernel1 = op1._func_table['kernel1'].root assert 'uvL0(t, x + 1, y + 1, z + 1) = make_float2(0,0);' in str(kernel1) # Check numerical correctness assert np.allclose(w.data, w1.data, rtol=1.e-6) def test_eqns_order_immunity(self): grid = Grid(shape=(30, 30, 30)) u = TimeFunction(name='u', grid=grid) v = TimeFunction(name='v', grid=grid) u1 = u.func(name='u1') v1 = v.func(name='v1') eqns = [Eq(v.forward, v.dx + v.dy + 1.), Eq(u.forward, u.dx + u.dy + 2.)] op0 = Operator(eqns, opt=opt0()) op1 = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) op0.apply(time_M=2) op1.apply(time_M=2, u=u1, v=v1) # Check generated code -- in particular, the fact that no # packing/unpacking is performed assert 's_yz2[ty][tz] = queue2[0]' in str(op1) assert 's_yz2[ty + 4][tz] = uvL0(t0, x + 1, y + 5, z + 1)' # Check numerical correctness assert np.allclose(u.data, u1.data, rtol=1.e-6) assert np.allclose(v.data, v1.data, rtol=1.e-6) @pytest.mark.parametrize('layers', [None, NoLayers]) def test_avoid_if_save_buffers(self, layers): nt = 10 grid = Grid(shape=(40, 40)) u = TimeFunction(name='u', grid=grid) usave = TimeFunction(name='usave', grid=grid, save=nt, layers=layers) vsave = usave.func(name='vsave') eqns = [Eq(u.forward, u + 1.), Eq(usave, u), Eq(vsave, u)] op = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) op.apply(time_M=nt-1) # Check generated code assert not any(i in str(op) for i in ('float2', 'float3', 'float4')) @pytest.mark.parametrize('layers', [None, NoLayers]) def test_avoid_if_streaming_in(self, layers): nt = 10 grid = Grid(shape=(40, 40)) u = TimeFunction(name='u', grid=grid, time_order=0) usave = TimeFunction(name='usave', grid=grid, save=nt, layers=layers) vsave = usave.func(name='vsave') # Init usave/vsave eqns0 = [Eq(usave, 1), Eq(vsave, 1)] Operator(eqns0)(time_M=nt-1) # Gradient-like Operator eqns1 = [Eq(u.forward, u + usave + vsave)] op1 = Operator(eqns1, opt=opt6({'par-tile': (4, 4)})) # Check generated code assert not any(i in str(op1) for i in ('float2', 'float3', 'float4')) op1.apply(time_M=nt-1) assert np.all(u.data == (1 + 1)*nt) @pytest.mark.parallel(mode=[(4, 'diag2')]) def test_w_mpi_w_overrides(self, mode): grid = Grid(shape=(40, 40, 40)) u = TimeFunction(name='u', grid=grid, space_order=4) v = u.func(name='v') eqns = [Eq(u.forward, u.dy + u.dz + v.dy + v.dz + .001), Eq(v.forward, u.dy + u.dz + v.dy + v.dz + .002)] op0 = Operator(eqns, opt=opt4({'par-tile': (4, 4)})) op1 = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) assert 'float2' not in str(op0) assert 'float2' in str(op1) grid1 = Grid(shape=(45, 50, 60)) u1 = TimeFunction(name='u1', grid=grid1, space_order=4) u2 = u1.func(name='u2') v1 = u1.func(name='v1') v2 = u1.func(name='v2') op0.apply(time_M=2, u=u1, v=v1) op1.apply(time_M=2, u=u2, v=v2) assert np.allclose(u1.data, u2.data, rtol=1.e-5) assert np.allclose(v1.data, v2.data, rtol=1.e-5) @pytest.mark.parallel(mode=[(4, 'basic'), (4, 'diag2')]) def test_carry_over_halo_metadata(self, mode): grid = Grid(shape=(80, 80, 80)) u = TimeFunction(name='u', grid=grid, space_order=4) v = u.func(name='v') w = u.func(name='w') eqns = [Eq(u.forward, u + .001), Eq(v.forward, v + .002), Eq(w.forward, (u.forward.dy2 + u.forward.dz2 + v.forward.dy2 + v.forward.dz2 + .003))] op0 = Operator(eqns, opt=opt4({'par-tile': (4, 4)})) op1 = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) assert 'float2' not in str(op0) assert 'float2' in str(op1) u1 = u.func(name='u1') v1 = u.func(name='v1') w1 = u.func(name='w1') op0.apply(time_M=2) op1.apply(time_M=2, u=u1, v=v1, w=w1) assert np.allclose(u.data, u1.data, rtol=1.e-5) assert np.allclose(v.data, v1.data, rtol=1.e-5) assert np.allclose(w.data, w1.data, rtol=1.e-5) @pytest.mark.parallel(mode=[(4, 'diag2')]) def test_minimum_halo_exchange_v0(self, mode): grid = Grid(shape=(80, 80, 80)) u = TimeFunction(name='u', grid=grid, space_order=4) v = u.func(name='v') w = u.func(name='w') eqns = [Eq(u.forward, u + .001), Eq(v.forward, v + .002), Eq(w.forward, u.forward.laplace + v.forward.dy2 + .003)] op0 = Operator(eqns, opt=opt4({'par-tile': (4, 4)})) op1 = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) # Check generated code assert 'float2' in str(op1) assert len([i for i in op1._func_table if i.startswith('haloupdate')]) == 1 assert len([i for i in op1._func_table if i.startswith('halowait')]) == 1 msgs = [i for i in op1.parameters if isinstance(i, MPIMsg)] assert len(msgs[0].halos) == 6 # All dims/dirs due to the laplacian assert len(msgs[1].halos) == 2 # Only y due to the .dy2 assert msgs[1].halos[0].side == (CENTER, LEFT, CENTER) assert msgs[1].halos[1].side == (CENTER, RIGHT, CENTER) u1 = u.func(name='u1') v1 = u.func(name='v1') w1 = u.func(name='w1') op0.apply(time_M=2) op1.apply(time_M=2, u=u1, v=v1, w=w1) assert np.allclose(u.data, u1.data, rtol=1.e-5) assert np.allclose(v.data, v1.data, rtol=1.e-5) assert np.allclose(w.data, w1.data, rtol=1.e-5) @pytest.mark.parallel(mode=[(4, 'diag2')]) def test_minimum_halo_exchange_v1(self, mode): grid = Grid(shape=(80, 80, 80)) u = TimeFunction(name='u', grid=grid, space_order=4) v = u.func(name='v') h = u.func(name='h') w = u.func(name='w') eqns = [Eq(u.forward, u + .001), Eq(v.forward, v + .002), Eq(h.forward, h + .003), Eq(w.forward, u.forward.laplace + v.forward.dy2 + h.forward.dy2)] op0 = Operator(eqns, opt=opt4({'par-tile': (4, 4)})) op1 = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) assert all(i in op1._func_table for i in ['haloupdate0', 'haloupdate1']) msgs = [i for i in op1.parameters if isinstance(i, MPIMsg)] assert len(msgs[1].halos) == 6 # All dims/dirs due to the laplacian assert len(msgs[0].halos) == 2 # Only y due to the .dy2 assert msgs[0].halos[0].side == (CENTER, LEFT, CENTER) assert msgs[0].halos[1].side == (CENTER, RIGHT, CENTER) u1 = u.func(name='u1') v1 = u.func(name='v1') h1 = u.func(name='h1') w1 = u.func(name='w1') op0.apply(time_M=2) op1.apply(time_M=2, u=u1, v=v1, h=h1, w=w1) assert np.allclose(w.data, w1.data, rtol=1.e-5) @pytest.mark.parallel(mode=[(1, 'diag2')]) def test_elastic_like_misc_codegen_mpi(self, mode): grid = Grid(shape=(10, 10, 10)) v = VectorTimeFunction(name='v', grid=grid, space_order=4) tau = TensorTimeFunction(name='tau', grid=grid, space_order=4) eqns = [Eq(v.forward, v - div(tau)), Eq(tau.forward, tau - diag(div(v.forward)))] op = Operator(eqns, opt=opt6()) op.cfunction # Check object abstraction by the IET's engine behaves as expected # (we used to have a spurious alias `f1` among op.parameters) op.arguments(time_M=1) assert all(i in str(op) for i in ('float2', 'float4')) # Check there are no redundant bundles bundlers = [i for i in op._func_table if i.startswith('bundle')] unbundlers = [i for i in op._func_table if i.startswith('unbundle')] assert len(bundlers) == len(unbundlers) == 3 # We used to generate loop blocking via rcompile for bundling/unbundling, # which is pointless, so we check we don't anymore funcs = [v for i, v in op._func_table.items() if i in bundlers + unbundlers] assert len(funcs) == 6 for i in funcs: assert len(FindNodes(Iteration).visit(i.root)) == 4 # Ensure that no redundant MPI halo exchanges are generated msgs = [i for i in op.parameters if isinstance(i, MPIMsg)] assert len(msgs) == 9 expected = [(2, tau[0]), (4, tau[1]), (2, tau[4]), (2, tau[8]), (4, tau[6]), (4, tau[5]), (2, v[0]), (2, v[1]), (2, v[2])] for msg, (exp_halos, exp_comp) in zip(msgs, expected): assert len(msg.halos) == exp_halos assert len(msg.target.components) == 1 assert msg.target.components[0] == exp_comp # There only are one variant of haloupdate and one variant of halowait # around, since the IET engine is expected to produce a single abstract # version of them, which works across vector types assert len([i for i in op._func_table if i.startswith('haloupdate')]) == 1 assert len([i for i in op._func_table if i.startswith('halowait')]) == 1 def test_avoid_make_float_with_reductions(self): grid = Grid(shape=(10, 10, 10)) u = TimeFunction(name='u', grid=grid, space_order=4) f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) eqns = [Eq(u.forward, u.dx), Inc(f, u.forward + 2.), Inc(g, u.forward + 3.)] op = Operator(eqns, opt=opt6({'par-tile': (4, 4)})) op.apply(time_M=9) assert np.all(f.data == 20) assert np.all(g.data == 30) @pytest.mark.parallel(mode=4) def test_acoustic_like_w_ricker_source(self, mode): space_order = 8 spacing = (20., 20., 20) origin = (0., 0., 0.) dt = 3.283 time_range = TimeAxis(start=0, stop=6000, step=dt) shape = (181 - space_order*2 - 40, 201 - space_order*2 - 40, 151 - space_order*2 - 40) vp_data = np.ones(shape, dtype=np.float32) * 1.5 vp_data[:, :, 40:] = 2.0 model = Model(vp=vp_data, origin=origin, shape=shape, nbl=0, bcs='damp', spacing=spacing, space_order=space_order) grid = model.grid vp = model.vp s = grid.stepping_dim.spacing p = TimeFunction(name='p', grid=grid, space_order=space_order, time_order=2) q = TimeFunction(name='q', grid=grid, space_order=space_order, time_order=2) p1 = p.func(name='p1') q1 = q.func(name='q1') src_coordinates = np.empty((1, 3)) src_coordinates[0, :] = np.array(model.domain_size) * .4 src_coordinates[:, 2] = spacing[2]*3.5 src = RickerSource(name='src', grid=grid, f0=0.01, a=1, t0=0, hicks_r=3, npoint=1, time_range=time_range, coordinates=src_coordinates) eqns = [Eq(p.forward, p.laplace * s**2 + p - p.backward + .4), Eq(q.forward, q.laplace * s**2 + q - q.backward + .4)] eqns += src.inject(field=p.forward, expr=src * s**2 * vp**2) eqns += src.inject(field=q.forward, expr=src * s**2 * vp**2) op0 = Operator(eqns, opt=opt4({'par-tile': (32, 8)})) op1 = Operator(eqns, opt=opt6({'par-tile': (32, 8)})) op0.apply(dt=dt, time_M=60) op1.apply(dt=dt, time_M=60, p=p1, q=q1) assert np.isclose(norm(p), norm(p1), rtol=1.e-5) assert np.isclose(norm(q), norm(q1), rtol=1.e-5) def test_issue_781(self): grid = Grid(shape=(32, 41, 21)) f = Function(name='f', grid=grid, space_order=8, avg_mode='safe_harmonic') g = Function(name='g', grid=grid, space_order=8) v = TimeFunction(name='v', grid=grid, save=Buffer(1), space_order=8, time_order=1) p = TensorTimeFunction(name='p', grid=grid, save=Buffer(1), space_order=8, time_order=1) dv = p.func(name='dv') eqns = [Eq(dv.forward[2], v.dz2), Eq(p.forward[0], g*p[0] - dv[0] - f), Eq(p.forward[2], g*p[2] - f), Eq(p.forward[3], g*p[1] - f), Eq(p.forward[5], g*p[5] - f)] op = Operator(eqns, opt=opt6({'par-tile': (16, 8, 8)})) op.cfunction class TestKernOptMisc: def test_minimize_symbols(self): grid = Grid(shape=(16, 16)) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=4, save=Buffer(2)) eqn = Eq(u.forward, u.dx + u.backward) op = Operator(eqn) op.cfunction kernel = op._func_table['kernel0'].root mds = [i for i in kernel.parameters if isinstance(i, ModuloDimension)] assert len(mds) == 2 class TestTemplateKernels: def _assert_jit(self, op): # NOTE: can't use .cfunction because we can't attempt to dynamically # load a templated Operator. But at least the following line will ensure # that the code compiles correctly op._compiler.jit_compile(op._soname, str(op)) def test_basic(self): grid = Grid(shape=(8, 8, 8)) u = TimeFunction(name='u', grid=grid) op = Operator(Eq(u.forward, u + 1), opt=('advanced', {'par-tile': None})) assert len(op.templates) == 3 def test_2p5_blocking(self): grid = Grid(shape=(8, 8, 8)) x, _, _ = grid.dimensions u = TimeFunction(name='u', grid=grid, space_order=4) op = Operator(Eq(u.forward, u.dx + 1), opt=opt1({'par-tile': (None, None)})) assert len(op.templates) == 2 trees = retrieve_iteration_tree(op._func_table['kernel0'].root) assert len(trees) == 2 assert len(trees[0]) == 1 and trees[0].root.dim is x assert trees[1].root.dim is x def test_kernelopt_default(self): grid = Grid(shape=(8, 8, 8)) f = TimeFunction(name='f', grid=grid, space_order=16) g = TimeFunction(name='g', grid=grid, space_order=16) eqn = Eq(f.forward, f + g.dz + g.dx + g.dy + 1) op = Operator(eqn, opt=opt4({'par-tile': (None, None)})) self._assert_jit(op) def test_kernelopt_default_v2(self): nstc = 8 grid = Grid(shape=(10, 10, 10)) _, _, z = grid.dimensions zo = Constant(name='zo', dtype=int) u = Function(name='u', grid=grid, space_order=(2*nstc, 0, 0), coefficients='symbolic') v = Function(name='v', grid=grid, space_order=(2*nstc, 0, 0)) f = Function(name='f', space_order=0, shape=(10,), dimensions=(z,)) f = f.evaluate.subs({z: z - zo}) eqn = Eq(v, f + u.dz2) op = Operator(eqn, opt=opt4({'par-tile': (None, None)})) self._assert_jit(op) def test_custom_coeffs_mixed_derivs(self): p = Dimension('p') nstc = 8 grid = Grid(shape=(10, 10, 10)) x, y, z = grid.dimensions f = Function(name='f', grid=grid, space_order=(2*nstc, 0, 0), coefficients='symbolic') g = Function(name='g', grid=grid, space_order=(2*nstc, 0, 0)) ax = Function(name='DD2x', space_order=0, shape=(2*nstc + 1,), dimensions=(p,)) ay = Function(name='DD2y', space_order=0, shape=(2*nstc + 1,), dimensions=(p,)) stencil_coeffs_x_p1 = Coefficient(1, f, x, ax) stencil_coeffs_y_p1 = Coefficient(1, f, y, ay) stencil_coeffs = Substitutions(stencil_coeffs_x_p1, stencil_coeffs_y_p1) eqn = Eq(g, f.dxdy, coefficients=stencil_coeffs) op = Operator(eqn, opt=opt4({'par-tile': (None, None), 'deriv-unroll': 'full'})) self._assert_jit(op) @skipif('nocuda') class TestAsyncPipeline: def _check_gencode(self, efunc, ncommits, nwaits, nasyncs, ngas=None): calls = FindNodes(Call).visit(efunc) for call, count in [('__pipeline_commit', ncommits), ('__pipeline_wait_prior', nwaits), ('__pipeline_memcpy_async', nasyncs)]: assert len([i for i in calls if i.name == call]) == count if ngas is not None: f, exp_ngas = ngas indexeds = FindSymbols('indexeds').visit(efunc) found = {i for i in indexeds if i.function is f} assert len(found) == exp_ngas @switchconfig(ignore_unknowns=True) @switchconfig(platform='ampere') @pytest.mark.parametrize('exprs,ncommits,nwaits,nasyncs,exact', [ # Simple cases (['u.dy + f'], 2, 1, 6, True), (['u.dy + f', 'v.dy + f'], 2, 1, 6, True), (['u.dy + u.dx + f'], 2, 1, 4, False), (['u.dy + u.dy.dx + f'], 4, 2, 12, True), (['v.dy + f', 'u.dy + f'], 2, 1, 6, True), (['u.dy.dx + u.dy.dz'], 4, 2, 12, True), (['u.dx.dz + f*w.dy.dz'], 2, 1, 6, True), (['u.dz + f*w.dy.dz + g*u.dy'], 2, 1, 12, True), (['u.dy + g*v.dy.dx'], 4, 2, 18, True), (['v.dy + v.dy.dx + f', 'u.dy + u.dy.dx + f'], 4, 2, 12, True), (['u.dx + g*v.dy.dx + f*w.dy.dz'], 4, 2, 18, True), (['u.laplace'], 2, 1, 8, True), # A-la F2F cases (['u.dy.dx + f*u.dx + f*u.dz.dx'], 4, 2, 10, True), (['u.dy.dx + u.dy.dz + f*u.dx + f*u.dz + u.dz2 + f*u.dz.dx'], 4, 2, 10, 'issue#499-2'), (['u.dx + u.dx.dz + u.dx.dy', 'f*u.dy.dx + g*u.dy.dz'], 4, 2, 10, True), (['u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz'], 4, 2, 10, 'issue#499-2'), (['u.dy.dx + u.dy.dz', 'v.dy.dx + v.dy.dz'], 4, 2, 12, False), (['u.dy.dx + u.dy.dz + f*u.dx.dz', 'v.dy.dx + v.dy.dz + f*v.dx.dz'], 4, 2, 10, False), (['u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz', 'v.dx2 + v.dy2 + v.dz2 + v.dy.dx + v.dy.dz + f*v.dx.dz'], 4, 2, 10, 'issue#499-2'), # More tricky cases (['u.dx.dz + g*u.dy'], 2, 1, 4, True), (['u.dx.dz + f*w.dy.dz + g*u.dy'], 2, 1, 10, True), (['u.dx + f*w.dy.dz.dx + g*u.dy.dz'], 4, 2, 16, True), (['u.dx.dz + u.dx.dy'], 4, 2, 10, True), (['u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy'], 4, 2, 20, 'issue#499'), (['u.dx + u.dx.dz + u.dx.dy', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz'], 4, 2, 22, 'issue#499'), (['u.dx + u.dx.dy + g*v.dy.dx', 'g*u.dy.dz + g*v.dx + u.dz'], 4, 2, 10, True), # A-la TTI Elastic (['u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy', 'u.dz.dx + f*v.dy.dz.dx + u.dz.dy + g*w.dx + u.dz'], 4, 2, 20, 'issue#499'), (['u.dx + u.dx.dz + u.dx.dy + g*v.dy.dx + f*w.dz.dx', 'f*u.dy.dx + f*w.dy.dz.dx + g*u.dy.dz + g*v.dx + f*u.dy', 'u.dz.dx + f*v.dy.dz.dx + u.dz.dy + g*w.dx + u.dz', 'v.dy + f*u.dx.dy + v.dy.dz + v.dy.dx + g*w.dx.dy', 'v.dz.dy + f*u.dx.dz.dy + v.dz.dx + g*w.dy + v.dz', 'w.dz + f*u.dx.dz + g*v.dy.dz + w.dz.dy + w.dz.dx'], 4, 2, 20, False), ]) @pytest.mark.parametrize('distance', [1, 2, 5]) def test_basic(self, exprs, ncommits, nwaits, nasyncs, exact, distance): grid = Grid(shape=(31, 17, 25)) x, _, _ = grid.dimensions nt = 5 f = Function(name='f', grid=grid, space_order=4) g = Function(name='g', grid=grid, space_order=4) functions = [] for name in ('u', 'v', 'w', 'a', 'b', 'c'): f0 = TimeFunction(name=name, grid=grid, space_order=4, save=nt, layers=NoLayers) f1 = f0.func(name='%s1' % name) functions.extend([f0, f1]) u, u1, v, v1, w, w1, a, a1, b, b1, c, c1 = functions fields = [u, v, w, a, b, c] eqns = [] for expr, field in zip(exprs, fields): eqns.append(Eq(field.forward, eval(expr) + 1.)) # NOTE: by using `nofloat3` we ensure that the order of arithmetic # operations in `op0` and `op1` is the same op0 = Operator(eqns, opt=opt7({'gpu-fit': fields, 'par-tile': (16, 4, 4), 'gpu-vect': 'nofloat3'})) op1 = Operator(eqns, opt=optX({'gpu-fit': fields, 'par-tile': (16, 4, 4), 'gpu-prefetch': f'1:{distance}'})) # Check generated code assert '__pipeline_memcpy_async' not in str(op0) assert ('__pipeline_wait_prior(0);' in str(op1) or '__pipeline_wait_prior(1);' in str(op1)) assert '__pipeline_memcpy_async' in str(op1) efunc = op1._func_table['kernel0'].root for i in FindNodes(Iteration).visit(efunc): if x in i.dim._defines: assert '__pipeline_wait_prior(1)' in str(i) if distance == 1: self._check_gencode(op1._func_table['kernel0'].root, ncommits, nwaits, nasyncs) if ncommits == nwaits: assert not any(i in str(op1) for i in ['comm0', 'comp0']) g.data[:] = 0.31 f.data[:] = 0.75 for i in (u, u1, v, v1, w, w1): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1, v=v1, w=w1) for f0, f1 in [(u, u1), (v, v1), (w, w1)]: if exact is True: # `gpu-prefetch` didn't alter the order of arithmetic operations! assert np.all(f0.data == f1.data) elif exact == 'issue#499': # `gpu-prefetch` didn't alter the order of arithmetic operations! assert np.all(f0.data[:nt-1] == f1.data[:nt-1]) elif exact == 'issue#499-2': assert np.allclose(f0.data[:nt-1], f1.data[:nt-1], rtol=1.e-3) else: assert np.allclose(f0.data, f1.data, rtol=1.e-5) @switchconfig(platform='ampere') @pytest.mark.parametrize('expr,expected', [ ('u.dx + f*u.dy.dx', 3), ]) def test_prefetch_both_core_and_halo(self, expr, expected): grid = Grid(shape=(42, 21, 33)) nt = 5 f = Function(name='f', grid=grid, space_order=4) # noqa u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = u.func(name='u1') eqn = Eq(u.forward, eval(expr)) op0 = Operator(eqn, opt=opt7({'gpu-fit': [u], 'par-tile': (16, 4, 4), 'gpu-vect': 'nofloat3'})) op1 = Operator(eqn, opt=optX({'gpu-fit': [u], 'par-tile': (16, 4, 4)})) op2 = Operator(eqn, opt=optX({'gpu-fit': [u], 'par-tile': (16, 4, 4), 'gpu-prefetch': 4})) # Check generated code self._check_gencode(op1._func_table['kernel0'].root, 4, 2, 10) self._check_gencode(op2._func_table['kernel0'].root, 4, 2, 14) f.data[:] = 0.75 for i in (u, u1): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) # NOTE: Expected exact equality since the only difference is the use of # `gpu-prefetch`, which doesn't alter order of arithmetic operations assert np.all(u.data == u1.data) # Operator options for SubDimension prefetch tests sd_opt = ('advanced', {'errctl': 'basic', 'gpu-opt': True, 'gpu-prefetch': (True,), 'gpu-glb-mem': 'async-pipe', 'par-tile': ((32, 4),)}) @switchconfig(platform='ampere') def test_issue_789(self): class SD(SubDomain): def define(self, dimensions): return {d: ('middle', 0, 0) for d in dimensions} grid = Grid(shape=(100, 100, 100)) sd = SD(grid=grid) v = TimeFunction(name="v", grid=grid) eq = Eq(v.forward, v, subdomain=sd) op = Operator(eq, opt=self.sd_opt) _ = op.ccode @switchconfig(platform='ampere') @pytest.mark.parametrize('side', ['left', 'middle', 'right']) def test_prefetch_subdimension(self, side): class SD(SubDomain): def define(self, dimensions): retval = {d: d for d in dimensions} if side == 'middle': val = ('middle', 5, 5) else: val = (side, 10) retval[dimensions[0]] = val return retval grid = Grid(shape=(20, 20, 20), extent=(19., 19., 19.)) sd = SD(grid=grid) v = TimeFunction(name="v", grid=grid) v.data_with_halo[0, :] = \ np.arange(v.shape_with_halo[1])[np.newaxis, :, np.newaxis] eq = Eq(v.forward, v.dy, subdomain=sd) op = Operator(eq, opt=self.sd_opt) # Check bounds if side == 'left': assert "__pipeline_memcpy_async(&s_yz0[(x_m & 1)]" in str(op.ccode) assert "for (int x = x_m; x <= x_m + x_ltkn0 - 1; x += 1)" in str(op.ccode) assert "if (x + 1 <= x_m + x_ltkn0 - 1)" in str(op.ccode) elif side == 'middle': assert "__pipeline_memcpy_async(&s_yz0[((x_m + x_ltkn0) & 1)]" \ in str(op.ccode) assert "for (int x = x_m + x_ltkn0; x <= x_M - x_rtkn0; x += 1)" \ in str(op.ccode) assert "if (x + 1 <= x_M - x_rtkn0)" in str(op.ccode) elif side == 'right': assert "__pipeline_memcpy_async(&s_yz0[((x_M - x_rtkn0 + 1) & 1)]" \ in str(op.ccode) assert "for (int x = x_M - x_rtkn0 + 1; x <= x_M; x += 1)" in str(op.ccode) assert "if (x + 1 <= x_M)" in str(op.ccode) # Check correctness op.apply(t_M=1) if side == 'left': assert np.all(np.isclose(v.data[1, :10], 1)) elif side == 'middle': assert np.all(np.isclose(v.data[1, 5:-5], 1)) elif side == 'right': assert np.all(np.isclose(v.data[1, -10:], 1)) @switchconfig(platform='ampere') def test_prefetch_subdimension_vel_like(self): class SD(SubDomain): def define(self, dimensions): return {d: ('middle', 0, 0) for d in dimensions} so = 8 grid = Grid(shape=(100, 100, 100)) x = grid.dimensions[0] sd = SD(grid=grid) b = Function(name='b', grid=grid, space_order=so) txx = TimeFunction(name="txx", grid=grid, time_order=1, space_order=so) txy = TimeFunction(name="txy", grid=grid, time_order=1, space_order=so) txz = TimeFunction(name="txz", grid=grid, time_order=1, space_order=so) vx = TimeFunction(name="vx", grid=grid, time_order=1, space_order=so) eq = Eq(vx.forward, vx + (txx.dx + txy.dy + txz.dz) * Derivative(b, (x, 0), x0={x: x+x.spacing/2}), subdomain=sd) op = Operator(eq, opt=self.sd_opt) assert "__pipeline_memcpy_async(&s_yz0[((x_m + x_ltkn0) & 1)]" in str(op.ccode) assert "for (int x = x_m + x_ltkn0; x <= x_M - x_rtkn0; x += 1)" in str(op.ccode) assert "if (x + 1 <= x_M - x_rtkn0)" in str(op.ccode) @switchconfig(platform='ampere') def test_wait_after_commit_w_smartsort(self): grid = Grid(shape=(42, 21, 33)) u = TimeFunction(name='u', grid=grid, space_order=4) eqn = Eq(u.forward, u.dx + u.dy.dx + 1.) op = Operator(eqn, opt=optX({'gpu-fit': [u], 'par-tile': (16, 4, 4), 'cse-algo': 'smartsort'})) # Check generated code calls = FindNodes(Call).visit(op._func_table['kernel0'].root) assert len(calls) == 21 # qq loop assert calls[6].name == '__pipeline_commit' assert calls[7].name == '__pipeline_wait_prior' # x loop assert calls[18].name == '__pipeline_commit' assert calls[19].name == '__pipeline_wait_prior' @switchconfig(platform='ampere') @pytest.mark.parametrize('axis', [0, 1, 2]) @pytest.mark.parametrize('prefetch', [False, True, 2, 3, 4]) @pytest.mark.parametrize('access', ['field', 'derivative', 'queue-front']) def test_noprefetch_1d_profile(self, axis, prefetch, access): """Cache profiles used as fields, derivatives and register queue fronts.""" grid = Grid(shape=(32, 16, 16)) d = grid.dimensions[axis] r = Function(name='r', grid=grid, shape=(grid.shape[axis],), dimensions=(d,), space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) exprs = { 'field': r*u.dz + u.dx, 'derivative': r.diff(d)*u + u.dx + u.dz, 'queue-front': (r*u).diff(d) + u.dx + u.dz, } equation = Eq(v.forward, exprs[access], subdomain=grid.interior) op = Operator(equation, name='cached_1d_profile', opt=optX({ 'cire-mode': '1pass', 'cire-hoist': 0, 'fission': False, 'gpu-vect': 'stream-only', 'gpu-prefetch': prefetch, 'par-tile': (16, 8, 8), })) kernels = [v.root for k, v in op._func_table.items() if k.startswith('kernel')] indexeds = [i for kernel in kernels for i in FindSymbols('indexeds').visit(kernel)] assert any(i.function is r for i in indexeds) # Async transfers are Calls; ordinary shared-memory loads are Expressions staged = [i for kernel in kernels for i in FindNodes((Expression, Call)).visit(kernel) if r in i.functions and any(w._mem_shared for w in i.writes)] assert not staged @switchconfig(platform='ampere') @pytest.mark.parametrize('ndim', [2, 3]) @pytest.mark.parametrize('prefetch', [False, 2, 4]) def test_shm_multidimensional_fields(self, ndim, prefetch): grid = Grid(shape=(32, 16, 16)[:ndim]) d = grid.dimensions[-1] f = Function(name='f', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) equation = Eq(v.forward, f.diff(d) + u.dx + u.diff(d)) op = Operator(equation, name='staged_multidimensional_field', opt=optX({ 'cire-mode': '1pass', 'cire-hoist': 0, 'fission': False, 'gpu-vect': 'stream-only', 'gpu-prefetch': prefetch, 'par-tile': (16, 8, 8)[:ndim], })) kernels = [v.root for k, v in op._func_table.items() if k.startswith('kernel')] staged = [i for kernel in kernels for i in FindNodes((Expression, Call)).visit(kernel) if f in i.functions and any(w._mem_shared for w in i.writes)] assert staged if ndim == 3 and prefetch: assert any(i.is_Call and i.name == '__pipeline_memcpy_async' for i in staged) @switchconfig(ignore_unknowns=True) @switchconfig(platform='ampere') @pytest.mark.parametrize('expr,ncommits,nwaits,nasyncs', [ ('f*u.dx + g*u.dy.dx', 4, 2, 12), ('v.dy + f + u ', 2, 1, 10), ]) def test_prefetch_isolated_functions(self, expr, ncommits, nwaits, nasyncs): grid = Grid(shape=(42, 21, 33)) nt = 5 f = Function(name='f', grid=grid, space_order=4) g = f.func(name='g') u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) v = TimeFunction(name='v', grid=grid, space_order=4, save=nt, layers=NoLayers) eqn = Eq(u.forward, eval(expr)) op0 = Operator(eqn, opt=opt7({'gpu-fit': [u], 'par-tile': (16, 4, 4)})) op1 = Operator(eqn, opt=optX({'gpu-fit': [u], 'par-tile': (16, 4, 4), 'gpu-prefetch': 2})) # Check generated code self._check_gencode( op1._func_table['kernel0'].root, ncommits, nwaits, nasyncs ) u1 = TimeFunction(name='u1', grid=grid, space_order=4, save=nt, layers=NoLayers) f.data[:] = 0.75 g.data[:] = 0.31 for i in (u, u1, v): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) # NOTE: Expected exact equality since the only difference is the use of # `gpu-prefetch`, which doesn't alter order of arithmetic operations assert np.all(u.data == u1.data) @switchconfig(platform='ampere') def test_prefetch_model_params_like_fallback(self): grid = Grid(shape=(40, 40, 40)) f = Function(name='f', grid=grid) s = SparseFunction(name='s', grid=grid, npoint=1, coordinates=[(0.5, 0.5, 0.5)]) eq = s.inject(field=f, expr=s + 1) # We just make sure it generates and compiles op = Operator(eq, opt=optX({'par-tile': (16, 4, 4), 'gpu-prefetch': 2})) op.cfunction @switchconfig(platform='ampere') @pytest.mark.parametrize('expr,ncommits,nwaits,nasyncs', [ ('u.dx + u.dy', 4, 2, 8), ('u.dx + f*u.dy', 4, 2, 10), ('u.dy + u.backward + f', 2, 1, 10), ]) def test_prefetch_uncentered(self, expr, ncommits, nwaits, nasyncs): grid = Grid(shape=(32, 40, 23)) nt = 5 f = Function(name='f', grid=grid, space_order=4) u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = u.func(name='u1') eqn = Eq(u.forward, eval(expr)) op0 = Operator(eqn, opt=opt7({'gpu-fit': [u], 'par-tile': (16, 4, 4)})) op1 = Operator(eqn, opt=optX({'gpu-fit': [u], 'par-tile': (16, 4, 4), 'gpu-prefetch': 4})) # Check generated code self._check_gencode(op1._func_table['kernel0'].root, ncommits, nwaits, nasyncs) f.data[:] = 0.75 for i in (u, u1): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) # NOTE: Expected exact equality since the only difference is the use of # `gpu-prefetch`, which doesn't alter order of arithmetic operations assert np.all(u.data == u1.data) @switchconfig(platform='ampere') def test_noprefetch_inplace_field(self): grid = Grid(shape=(16, 8, 8)) hess = Function(name='hess', grid=grid) op = Operator(Eq(hess, hess + 2), name='noprefetch_inplace_field', opt=optX({'par-tile': (8, 4, 4), 'gpu-prefetch': 2})) kernel = op._func_table['kernel0'].root self._check_gencode(kernel, 0, 0, 0) hess.data[:] = 1 op.apply() assert np.all(hess.data == 3) @switchconfig(platform='ampere') @pytest.mark.parametrize('nslots', [1, 2]) def test_prefetch_buffered_field(self, nslots): grid = Grid(shape=(53, 17, 37)) v = TimeFunction(name='v', grid=grid, space_order=16, save=Buffer(nslots)) tau = TimeFunction(name='tau', grid=grid, space_order=16) v1 = v.func(name='v1') eqn = Eq(v.forward, v + tau.dx + tau.dy + tau.dz) opt = {'par-tile': (32, 8, 48), 'gpu-vect': 'stream-minimal', 'cire-hoist': 0} op0 = Operator(eqn, name='buffered_field_reference', opt=optX({**opt, 'gpu-prefetch': False})) op1 = Operator(eqn, name='prefetch_buffered_field', opt=optX({**opt, 'gpu-prefetch': 4})) kernel = op1._func_table['kernel0'].root exprs = FindNodes(Expression).visit(kernel) assert not [i for i in exprs if v in i.reads] assert len([i for i in exprs if i.write is v]) == 1 calls = FindNodes(Call).visit(kernel) loads = [i for i in calls if i.name == '__pipeline_memcpy_async' and v in FindSymbols().visit(i)] assert len(loads) == 2 for i in (v, v1, tau): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, v=v1) assert np.allclose(v.data, v1.data, rtol=1.e-5) def run_etti_stress_like_incremental(self, exprs, ncommits, nwaits, nasyncs, ngas, extra, tol=None): grid = Grid(shape=(31, 17, 25)) x, y, z = grid.dimensions nt = 5 # These are necessary to prevent derivative factorization f = Function(name='f', grid=grid, space_order=16) g = Function(name='g', grid=grid, space_order=16) v = VectorTimeFunction(name="v", grid=grid, space_order=16) tau = TensorTimeFunction(name="tau", grid=grid, space_order=16, save=nt, layers=NoLayers) vx, vy, vz = v fields = tau.values() C = .001*sum(fields) # noqa exprs, eqns = as_tuple(exprs), [] for expr, field in zip(exprs, fields): eqns.append(Eq(field.forward, eval(expr) + 1.)) cire_mode = extra.get('cire-mode', '1pass') op0 = Operator(eqns, opt=optX({'gpu-fit': fields, 'par-tile': (16, 8), 'gpu-vect': False, # Legacy behaviour: 'cire-hoist': 0, 'cire-mode': cire_mode, 'fission': 'aggressive:200'})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'gpu-fit': fields, 'par-tile': (16, 8), 'gpu-vect': False, # Legacy behaviour: 'cire-hoist': 0, 'cire-mode': cire_mode, 'fission': 'aggressive:200', **extra})) # Check generated code if extra['gpu-prefetch'] in ('1:1', '4:1:1', '4:1:4'): if ngas is not None: kernel = op1._func_table['kernel0'].root self._check_gencode( kernel, ncommits, nwaits, nasyncs, (eval(ngas[0]), ngas[1]) ) else: assert ncommits == 0 f.data[:] = 19.0 g.data[:] = 18.0 vx.data[:] = 11.0 vy.data[:] = 12.0 vz.data[:] = 7.0 overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} for objs in (fields, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) if tol is None: tol = 1.e-3 for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=tol) return op0, op1 @switchconfig(ignore_unknowns=True) @switchconfig(platform='ampere') @pytest.mark.parametrize('exprs,ncommits,nwaits,nasyncs,ngas', [ (('g*vz.dx+f*vz.dy', 'g*vz.dx+f*vz.dy'), 6, 3, 34, ('vz', 19)), (('vz.dz+g*vz.dx+f*vz.dy', 'vz.dz+g*vz.dx+f*vz.dy'), 6, 3, 34, ('vz', 19)), (('vz.dx', 'vz.dx', 'vz.dx', 'vz.dy', 'vz.dy', 'vz.dz'), 6, 3, 14, ('vz', 15)), (('vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + fields[0] + 1', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + fields[1] + 1', 'vx.dx + f*vy.dx + vx.dy + g*vz.dx + vx.dz + fields[2] + 1', 'vy.dy + f*vx.dy + vy.dz + vy.dx + g*vz.dy + fields[3] + 1', 'vy.dy + f*vx.dy + vy.dx + g*vz.dy + vy.dz + fields[4] + 1', 'vz.dz + f*vx.dz + g*vy.dz + vz.dy + vz.dx + fields[5] + 1'), 6, 3, 54, ('vz', 15)), ]) @pytest.mark.parametrize('extra', [ {'gpu-prefetch': '1:1'}, {'gpu-prefetch': '1:2'}, {'gpu-prefetch': '1:5'}, ]) def test_etti_stress_like_incremental(self, exprs, ncommits, nwaits, nasyncs, ngas, extra): self.run_etti_stress_like_incremental( exprs, ncommits, nwaits, nasyncs, ngas, extra ) @switchconfig(ignore_unknowns=True) @switchconfig(platform='ampere') @pytest.mark.parametrize('ths,ncommits,nwaits', [ (1, 20, 10), (4, 6, 3), ]) def test_etti_stress_like_thsX(self, ths, ncommits, nwaits): exprs = ('vx.dx + vx.dz + vx.dy + f*vy.dx + f*vz.dx + fields[0] + 1', 'vx.dx + vz.dx + vx.dz + g*vy.dx + vx.dy + fields[1] + 1', 'vx.dx + vy.dx + vx.dy + g*vz.dx + vx.dz + fields[2] + 1') # NOTE: only the ncommits and nwaits are expected to vary across different # thresholds; the number of async copies or global mem accesses is exactly # the same in all cases instead self.run_etti_stress_like_incremental( exprs, ncommits, nwaits, 48, ('vx', 16), {'gpu-prefetch': f'4:1:{ths}'}, tol=1.e-2 ) def run_etti_velocity_like(self, distance, extra, mode=1, ths=np.inf, expected=None): grid = Grid(shape=(42, 15, 31)) x, y, z = grid.dimensions nt = 5 f = Function(name='f', grid=grid, space_order=16) v = VectorTimeFunction(name="v", grid=grid, space_order=16) tau = TensorTimeFunction(name="tau", grid=grid, space_order=16) fields = tuple(v) expr = v.dt - f * div(tau) eqns = [Eq(v.forward, solve(expr, v.forward))] op0 = Operator(eqns, opt=optX({'gpu-fit': fields, 'par-tile': (16, 8), 'gpu-vect': False, 'gpu-prefetch': False})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'cire-hoist': 0, # Legacy 'gpu-fit': fields, 'par-tile': (16, 8), 'gpu-prefetch': f'{mode}:{distance}:{ths}', **extra})) if expected is not None: ncommits, nwaits, nasyncs, ngas = expected kernel0 = op1._func_table['kernel0'].root self._check_gencode(kernel0, ncommits, nwaits, nasyncs, (v[2], ngas)) f.data[:] = 3.5 init_chessboard(tau[0]) init_chessboard(tau[1]) init_chessboard(tau[2]) init_chessboard(tau[4]) init_chessboard(tau[5]) init_chessboard(tau[8]) overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} for objs in (fields, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2, dt=.001) op1.apply(time_M=nt-2, dt=.001, **overrides) for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) return op0, op1 @switchconfig(platform='ampere') @pytest.mark.parametrize('distance,expected,extra', [ (1, (2, 1, 34, 2), {'gpu-vect': False}), (2, (3, 1, 51, 2), {'gpu-vect': False}), ]) def test_etti_velocity_like_mode1(self, distance, expected, extra): self.run_etti_velocity_like(distance, extra, mode=1, expected=expected) @switchconfig(platform='ampere') @pytest.mark.parametrize('distance,ths,expected,extra', [ (1, 4, (4, 2, 56, 3), {'gpu-vect': False}), (2, 4, (6, 2, 84, 4), {'gpu-vect': False}), ]) def test_etti_velocity_like_mode4_thsX(self, distance, ths, expected, extra): self.run_etti_velocity_like( distance, extra, mode=4, ths=ths, expected=expected ) def test_etti_velocity_like_with_mu_ths4(self): grid = Grid(shape=(42, 15, 31)) x, y, z = grid.dimensions f = Function(name='f', grid=grid, space_order=8) m = Function(name='m', grid=grid, space_order=8, avg_mode='safe_harmonic') v = VectorTimeFunction(name="v", grid=grid, space_order=8) tau = TensorTimeFunction(name="tau", grid=grid, space_order=8) expr = v.dt - f * div(tau) eq = Eq(v.forward, solve(expr, v.forward)) mmask = INT(Gt(m, Cast(0, dtype=m.dtype))) sheartau = {t for t in tau.values() if t not in tau.diagonal()} eq = eq.subs({der: mmask * der for der in retrieve_derivatives(eq) if der.expr in sheartau}) op = Operator(eq, opt=optX({'par-tile': (32, 4, 16), 'cire-hoist': 0, # Legacy 'gpu-vect': False, 'gpu-prefetch': '4:1:4'})) op.cfunction # Check generated code kernel = op._func_table['kernel0'].root self._check_gencode(kernel, 4, 2, 58, (v[2], 3)) @switchconfig(platform='ampere') def test_shrunk_halo(self): grid = Grid(shape=(16, 16, 16)) # Shrinking the halo is crucial for this test as it will trigger # parts of the compiler that are expected to inform the runtime about # the safety of all the memory accesses so = (4, 2, 2) # NOTE: we need Vector/TensorTimeFunction, rather than plain TimeFunction, # to trigger interpolations v = VectorTimeFunction(name="v", grid=grid, space_order=so) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so) v_y = v[1] tau_xx = tau[0] eqn = Eq(tau_xx.forward, tau_xx + v_y.dx) op0 = Operator(eqn, opt=opt4({'errctl': 'basic', 'par-tile': (32, 4, 24)})) op1 = Operator(eqn, opt=optX({'errctl': 'basic', 'gpu-vect': False, 'par-tile': (32, 4, 24)})) # Check generated code self._check_gencode(op1._func_table['kernel0'].root, 4, 2, 8) tau_xx1 = tau_xx.func(name='tau_xx1') v_y.data[:] = 1.5 op0.apply(time_M=1) op1.apply(time_M=1, tau_xx=tau_xx1) assert np.all(tau_xx.data == tau_xx1.data) @switchconfig(platform='ampere') def test_multiple_prefetch_levels(self): grid = Grid(shape=(16, 16, 16)) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u + v.dy + 1), Eq(v.forward, u.forward.dx + v.dy + 1)] op = Operator(eqns, opt=optX({'par-tile': (16, 4, 4), 'gpu-prefetch': [2, False]})) # Check generated code self._check_gencode(op._func_table['kernel0'].root, 2, 1, 8) # Below, since gpu-glb-mem=async-pipe (via optX) and gpu-prefetch=False, # we expect batching only self._check_gencode(op._func_table['kernel1'].root, 1, 1, 3) @switchconfig(ignore_unknowns=True) @switchconfig(platform='ampere') def test_partile_smaller_than_halo(self): grid = Grid(shape=(61, 27, 45)) x, y, z = grid.dimensions nt = 5 v = VectorTimeFunction(name="v", grid=grid, space_order=16) tau = TensorTimeFunction(name="tau", grid=grid, space_order=16, save=nt, layers=NoLayers) vx, vy, vz = v fields = tau.values() exprs = [vx.dx + vx.dz + vx.dy + 1, vx.dx + vx.dz + vx.dy + 1, vx.dx + vx.dy + vx.dz + 1, vx.dy + 1, vx.dy + 1, vx.dz + 1] eqns = [Eq(field.forward, expr + 1.) for expr, field in zip(exprs, fields)] op0 = Operator(eqns, opt=optX({'gpu-fit': fields, 'par-tile': (32, 8), 'gpu-vect': False, 'gpu-prefetch': False})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'gpu-fit': fields, 'par-tile': (16, 4), 'deriv-uniformize': True, 'gpu-vect': False})) # Check generated code self._check_gencode(op1._func_table['kernel0'].root, 4, 2, 18) vx.data[:] = 11.0 vy.data[:] = 12.0 vz.data[:] = 7.0 overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} for objs in (fields, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) @switchconfig(platform='ampere') def test_derivative_of_nonatom(self): grid = Grid(shape=(40, 30, 20)) x, y, z = grid.dimensions space_order = 16 vx = TimeFunction(name='vx', grid=grid, space_order=space_order, time_order=1, save=Buffer(1)) vy = vx.func(name='vy') txx = vx.func(name='txx') txy = vx.func(name='txy') f0 = Function(name='f0', grid=grid, space_order=space_order) f1 = f0.func(name='f1') g = f0.func(name='g') eqns = [ Eq(txx.forward, (g * cos(f1) * cos(f0) * vx).dx), Eq(txy.forward, vy.dy + txy) ] op0 = Operator(eqns, subs=grid.spacing_map, opt=opt4({'par-tile': (16, 8)})) op1 = Operator(eqns, subs=grid.spacing_map, opt=optX({'errctl': 'basic', 'gpu-vect': False, 'gpu-prefetch': 4, 'par-tile': (16, 8)})) self._check_gencode(op1._func_table['kernel1'].root, 4, 2, 20) txx1 = txx.func(name='txx1') txy1 = txx.func(name='txy1') f0.data[:] = 0.25 f1.data[:] = 0.5 g.data[:] = 2.0 for i in (vx, vy, txx, txy, txx1, txy1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, txx=txx1, txy=txy1) # NOTE: Expected exact equality since the only difference is the use of # `gpu-prefetch`, which doesn't alter the order of arithmetic operations assert np.all(txx.data == txx1.data) assert np.all(txy.data == txy1.data) @switchconfig(platform='ampere') def test_two_pass_v0(self): grid = Grid(shape=(40, 40, 40)) u = TimeFunction(name='u', grid=grid, space_order=16) eqn = Eq(u.forward, u.dy.dz) op = Operator(eqn, opt=optX({'par-tile': (16, 8), 'cire-mode': '2pass'})) op.cfunction # Check generated code self._check_gencode(op._func_table['kernel0'].root, 2, 1, 6) self._check_gencode(op._func_table['kernel1'].root, 2, 1, 6) @switchconfig(platform='ampere') def test_two_pass_v1(self): grid = Grid(shape=(40, 40, 40)) x, y, z = grid.dimensions interp_order = 2 u = TimeFunction(name='u', grid=grid, space_order=8, time_order=1, save=Buffer(1)) v = u.func(name='v') y0m = y - y.spacing/2 x0p = x + x.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(interp_order, interp_order)) expr = Derivative(u.dx, x, y, x0={x: x0p, y: y0m}, **kw) eqn = Eq(v.forward, expr + v + .001) op0 = Operator(eqn, opt=opt4({'cire-mode': '2pass'})) op1 = Operator(eqn, opt=optX({'errctl': 'basic', 'par-tile': (32, 4), 'cire-mode': '2pass'})) v1 = v.func(name='v1') for i in (u, v, v1): init_chessboard(i) op0(time_M=5) op1(time_M=5, v=v1) assert np.allclose(v.data, v1.data, rtol=1.e-5) @switchconfig(platform='ampere') @pytest.mark.parametrize('pfc', ['1:1', '1:2', '4:1']) def test_two_pass_v2(self, pfc): grid = Grid(shape=(41, 35, 29)) x, y, z = grid.dimensions f = Function(name='f', grid=grid, space_order=8) vz = TimeFunction(name='vz', grid=grid, space_order=8, save=Buffer(1)) vy = vz.func(name='vy') tyz = vz.func(name='tyz') x0m = x - x.spacing/2 y0m = y - y.spacing/2 x0p = x + x.spacing/2 y0p = y + y.spacing/2 z0p = z + z.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(2, 2)) a = cos(f) b = sin(f) expr = ( Derivative(a*vy.forward, x, z, x0={x: x0m, z: z0p}, **kw).dx(x0=x0p) + Derivative(b*vy.forward, y, z, x0={y: y0m, z: z0p}, **kw).dy(x0=y0p) ) eqns = [Eq(vy.forward, tyz.dx(x0=x0m).dx.dz), Eq(tyz.forward, vz.dx.dy + expr + tyz)] op0 = Operator(eqns, opt=optX({'errctl': 'basic', 'cire-mode': ('1pass', '1pass', '2pass'), 'gpu-prefetch': False, 'gpu-vect': False, 'par-tile': (32, 4, 4)})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'cire-mode': ('1pass', '1pass', '2pass'), 'gpu-prefetch': [False, False, pfc, False], 'gpu-vect': False, 'par-tile': (32, 4, 4)})) vy1 = vy.func(name='vy1') tyz1 = tyz.func(name='tyz1') f.data[:] = 0.25 for i in (vy, vz, tyz, vy1, tyz1): init_chessboard(i) op0(dt=0.01, h_x=10, h_y=10, h_z=10, time_M=2) op1(dt=0.01, h_x=10, h_y=10, h_z=10, time_M=2, vy=vy1, tyz=tyz1) assert np.linalg.norm(vy.data) > 0 assert np.isfinite(np.linalg.norm(vy.data)) assert np.allclose(vy.data, vy1.data, rtol=1.e-5) assert np.allclose(tyz.data, tyz1.data, rtol=1.e-5) @pytest.mark.parametrize('distance', [1, 2, 5]) def test_2pass_v3(self, distance): grid = Grid(shape=(100, 37, 50)) so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so) v = u.func(name='v') eqns = [ Eq(u.forward, u.dy.dx + u.dy.dz + f*u.dx.dz), Eq(v.forward, v.dy.dx + v.dy.dz + f*v.dx.dz) ] op0 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': (32, 4, 48)})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'par-tile': (32, 4, 48), 'cire-mode': 'x|y', 'gpu-prefetch': f'4:{distance}', 'gpu-vect': False})) u1 = u.func(name='u1') v1 = v.func(name='v1') for i in (u, v, u1, v1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, u=u1, v=v1) assert np.allclose(u.data, u1.data, rtol=1.e-2) assert np.allclose(v.data, v1.data, rtol=1.e-2) @switchconfig(ignore_unknowns=True) @switchconfig(platform='ampere') @pytest.mark.parametrize('prefetch,ncommits,nwaits,nasyncs', [ (2, 0, 0, 0), (3, 4, 2, 4), ]) def test_2pass_v4(self, prefetch, ncommits, nwaits, nasyncs): grid = Grid(shape=(40, 40, 40)) u = TimeFunction(name='u', grid=grid, space_order=16) eqn = Eq(u.forward, u.dx) op = Operator(eqn, opt=optX({'par-tile': (16, 8), 'gpu-prefetch': prefetch, 'cire-mode': 'x'})) op.cfunction # Check generated code self._check_gencode(op._func_table['kernel0'].root, ncommits, nwaits, nasyncs) @switchconfig(platform='ampere') def test_prefetch_from_nonperfect_deriv(self): grid = Grid(shape=(40, 30, 20)) g = Function(name='g', grid=grid, space_order=8) vx = TimeFunction(name='vx', grid=grid, space_order=8, time_order=1, save=Buffer(1)) txx = vx.func(name='txx') eqn = Eq(txx.forward, (cos(g)*vx.forward).dx.dy + txx) op0 = Operator(eqn, subs=grid.spacing_map, opt=opt4({'par-tile': (16, 8)})) op1 = Operator(eqn, subs=grid.spacing_map, opt=optX({'errctl': 'basic', 'gpu-vect': False, 'gpu-prefetch': 4, 'par-tile': (16, 8)})) # Check generated code # Both `vx` and the temporary carrying `cos(g)` are expected to # be prefetched self._check_gencode(op1._func_table['kernel1'].root, 4, 2, 26) txx1 = txx.func(name='txx1') g.data[:] = 0.25 for i in (vx, txx, txx1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, txx=txx1) assert np.allclose(txx.data, txx1.data, rtol=1.e-5) @switchconfig(platform='ampere') def test_nonatomic_derivs(self): grid = Grid(shape=(40, 30, 20)) x, y, z = grid.dimensions space_order = 8 interp_order = 2 f = Function(name='f', grid=grid, space_order=space_order) g = f.func(name='g') vx = TimeFunction(name='vx', grid=grid, space_order=space_order, time_order=1, save=Buffer(1)) txx = vx.func(name='txx') x0m = x - x.spacing/2 y0m = y - y.spacing/2 y0p = y + y.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(interp_order, interp_order)) R = cos(g)*sin(f)*cos(f) expr = Derivative(R*vx, x, y, x0={x: x0m, y: y0p}, **kw).dy(x0=y0m) eqn = Eq(txx.forward, expr + txx) op = Operator(eqn, subs=grid.spacing_map, opt=optX({'errctl': 'basic', 'gpu-vect': False, 'par-tile': (16, 8)})) op.cfunction # Check generated code self._check_gencode(op._func_table['kernel1'].root, 4, 2, 24, (vx, 12)) @switchconfig(platform='ampere') def test_explicit_interp_complex(self): grid = Grid(shape=(40, 30, 20)) x, y, z = grid.dimensions space_order = 8 interp_order = 2 vx = TimeFunction(name='vx', grid=grid, space_order=space_order, time_order=1, save=Buffer(1)) vy = vx.func(name='vy') vz = vx.func(name='vz') txx = vx.func(name='txx') txy = vx.func(name='txy') txz = vx.func(name='txz') tyy = vx.func(name='tyy') tyz = vx.func(name='tyz') tzz = vx.func(name='tzz') f = Function(name='f', grid=grid, space_order=space_order) g = f.func(name='g') h0 = f.func(name='h0') h1 = f.func(name='h1') h2 = f.func(name='h2') R00 = cos(g) R01 = sin(g)*cos(f) R02 = sin(g)*sin(f) R10 = sin(f) R11 = cos(f) R12 = -cos(g)*sin(f) R20 = -sin(g)*cos(f) R21 = cos(g)*sin(f) R22 = cos(g) x0m, y0m, z0m = x - x.spacing/2, y - y.spacing/2, z - z.spacing/2 x0p, y0p, z0p = x + x.spacing/2, y + y.spacing/2, z + z.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(interp_order, interp_order)) def foo0(field): return ( (R00*field).dx(x0=x0m) + Derivative(R01*field, x, y, x0={x: x0m, y: y0p}, **kw).dy(x0=y0m) + Derivative(R02*field, x, z, x0={x: x0m, z: z0p}, **kw).dz(x0=z0m) ) def foo1(field): return ( Derivative(R10*field, x, y, x0={x: x0p, y: y0m}, **kw).dx(x0=x0m) + (R11*field).dy(x0=y0m) + Derivative(R12*field, y, z, x0={y: y0m, z: z0p}, **kw).dz(x0=z0m) ) def foo2(field): return ( Derivative(R20*field, x, z, x0={x: x0p, z: z0m}, **kw).dx(x0=x0m) + Derivative(R21*field, y, z, x0={y: y0p, z: z0m}, **kw).dy(x0=y0m) + (R22*field).dz(x0=z0m) ) def foo3(field): return ( (R00*field).dx(x0=x0p) + Derivative(R01*field, x, y, x0={x: x0p, y: y0m}, **kw).dy(x0=y0p) + Derivative(R02*field, x, z, x0={x: x0p, z: z0m}, **kw).dz(x0=z0p) ) def foo4(field): return ( Derivative(R10*field, x, y, x0={x: x0m, y: y0p}, **kw).dx(x0=x0p) + (R11*field).dy(x0=y0p) + Derivative(R12*field, y, z, x0={y: y0p, z: z0m}, **kw).dz(x0=z0p) ) def foo5(field): return ( Derivative(R20*field, x, z, x0={x: x0m, z: z0p}, **kw).dx(x0=x0p) + Derivative(R21*field, y, z, x0={y: y0m, z: z0p}, **kw).dy(x0=y0p) + (R22*field).dz(x0=z0p) ) eqns = [ Eq(txx.forward, (h0*foo0(vx) + h1*foo1(vy) + h2*foo2(vz)) + txx), Eq(tyy.forward, (h1*foo0(vx) + h1*foo1(vy) + h2*foo2(vz)) + tyy), Eq(tzz.forward, (h2*foo0(vx) + h1*foo1(vy) + h2*foo2(vz)) + tzz), Eq(tyz.forward, (h0*(foo5(vy) + foo4(vz))) + tyz), Eq(txz.forward, (h1*(foo5(vx) + foo3(vz))) + txz), Eq(txy.forward, (h2*(foo4(vx) + foo3(vy))) + txy), ] op = Operator(eqns, subs=grid.spacing_map, opt=optX({'gpu-vect': False, 'par-tile': (16, 8), 'deriv-collect': True})) op.cfunction # Check generated code TestKernelOptDefault()._check_kernel_opt(op, vx, 21, 12, 10, 15, exp_tmps=9, kids=1) TestKernelOptDefault()._check_kernel_opt(op, vx, 21, 10, 9, 13, exp_tmps=0, kids=2) @switchconfig(platform='ampere') def test_queue_heuristic0_revamped(self): grid = Grid(shape=(31, 17, 25)) nt = 5 # These are necessary to prevent derivative factorization f = Function(name='f', grid=grid, space_order=8) g = Function(name='g', grid=grid, space_order=8) v = VectorTimeFunction(name="v", grid=grid, space_order=8) tau = TensorTimeFunction(name="tau", grid=grid, space_order=8, save=nt, layers=NoLayers) _, _, vz = v.values() tauxx, tauxy, _, _, _, _ = fields = tau.values() eqns = [Eq(tauxx.forward, g*vz.dx + f*vz.dy + 1.), Eq(tauxy.forward, g*vz.dx + f*vz.dy + 1.)] op0 = Operator(eqns, opt=opt0({'errctl': 'basic', 'gpu-fit': fields, 'par-tile': (16, 8)})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'gpu-fit': fields, 'gpu-vect': False, 'par-tile': (16, 8)})) f.data[:] = 19.0 g.data[:] = 18.0 vz.data[:] = 7.0 tauxx1 = tauxx.func(name='tauxx1') tauxy1 = tauxy.func(name='tauxy1') tauxx.data[:] = 100 tauxy.data[:] = 200 tauxx1.data[:] = 100 tauxy1.data[:] = 200 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, tau_xx=tauxx1, tau_xy=tauxy1) assert np.allclose(tauxx.data, tauxx1.data, rtol=1.e-3) assert np.allclose(tauxy.data, tauxy1.data, rtol=1.e-3) @switchconfig(platform='ampere') @pytest.mark.parametrize('distance,ncommits,nwaits,nasyncs', [ (1, 6, 3, 26), (2, 9, 3, 39) ]) def test_prefetch_multi_phases(self, distance, ncommits, nwaits, nasyncs): grid = Grid(shape=(31, 17, 25)) nt = 5 f = Function(name='f', grid=grid, space_order=4) g = f.func(name='g') h = f.func(name='h') u = TimeFunction(name='u', grid=grid, space_order=4) u1 = u.func(name='u1') u2 = u.func(name='u2') v = u.func(name='v') eqn = Eq(u.forward, f*u.dy.dz + v.dy + h*u.dx.dz + g*f*v.dy.dx + 1) op0 = Operator(eqn, opt=opt7({'par-tile': (16, 4, 4), 'gpu-vect': False})) op1 = Operator(eqn, opt=optX({'par-tile': (16, 4, 4), 'gpu-vect': False, 'gpu-prefetch': f'4:{distance}'})) op2 = Operator(eqn, opt=optX({'par-tile': (16, 4, 4), 'gpu-vect': False, 'gpu-prefetch': f'4:{distance}:2'})) # Check generated code: ths=2 should split prefetching into multiple phases kernel = op2._func_table['kernel0'].root self._check_gencode(kernel, ncommits, nwaits, nasyncs) # Check numerical correctness f.data[:] = 0.75 g.data[:] = 0.31 h.data[:] = 0.25 init_chessboard(u) init_chessboard(u1) init_chessboard(u2) init_chessboard(v) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) op2.apply(time_M=nt-2, u=u2) assert np.allclose(u1.data, u2.data, rtol=1.e-4) assert np.allclose(u.data, u2.data, rtol=1.e-4) @switchconfig(platform='ampere') @pytest.mark.parametrize('distance,ncommits,nwaits,nasyncs', [ (1, 2, 1, 6), (2, 3, 1, 9), ]) def test_prefetch_multi_phases_fallback(self, distance, ncommits, nwaits, nasyncs): grid = Grid(shape=(31, 17, 25)) u = TimeFunction(name='u', grid=grid, space_order=4) eqn = Eq(u.forward, u.dy + 1) op0 = Operator(eqn, opt=optX({'par-tile': (16, 4, 4), 'gpu-vect': False, 'gpu-prefetch': f'4:{distance}'})) op1 = Operator(eqn, opt=optX({'par-tile': (16, 4, 4), 'gpu-vect': False, 'gpu-prefetch': f'4:{distance}:1'})) # Check generated code efunc4 = op0._func_table['kernel0'].root efunc5 = op1._func_table['kernel0'].root self._check_gencode(efunc4, ncommits, nwaits, nasyncs) self._check_gencode(efunc5, ncommits, nwaits, nasyncs) @switchconfig(platform='ampere') @pytest.mark.parametrize('exprs,ths,ncommits,nwaits,nasyncs', [ (('u.dy + g*v.dy.dx', 'v'), 1, 9, 3, 30), (('f*u.dy.dz + h*u.dx.dz + g*f*v.dy.dx', 'f*v.dy.dz + h*v.dx.dz + g*f*u.dy.dx'), 2, 9, 3, 45), ]) def test_ths_in_tandem_w_distance2(self, exprs, ths, ncommits, nwaits, nasyncs): grid = Grid(shape=(31, 17, 25)) nt = 5 f = Function(name='f', grid=grid, space_order=4) g = f.func(name='g') h = f.func(name='h') u = TimeFunction(name='u', grid=grid, space_order=4, layers=NoLayers) v = u.func(name='v') u1 = u.func(name='u1') v1 = v.func(name='v1') eqns = [Eq(u.forward, eval(exprs[0]) + 1), Eq(v.forward, eval(exprs[1]) + 2)] op0 = Operator(eqns, name='op0', opt=optX({'par-tile': (16, 4, 4), 'gpu-vect': False, 'gpu-prefetch': '4:2'})) op1 = Operator(eqns, name='op1', opt=optX({'par-tile': (16, 4, 4), 'gpu-vect': False, 'gpu-prefetch': f'4:2:{ths}'})) # Check generated code kernel = op1._func_table['kernel0'].root self._check_gencode(kernel, ncommits, nwaits, nasyncs) f.data[:] = 0.75 g.data[:] = 0.31 h.data[:] = 0.25 for i in (u, v, u1, v1): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1, v=v1) assert np.allclose(u.data, u1.data, rtol=1.e-7) assert np.allclose(v.data, v1.data, rtol=1.e-7) @pytest.mark.parallel(mode=[(1, 'diag2')]) def test_mpi(self, mode): grid = Grid(shape=(37, 45, 40)) u = TimeFunction(name='u', grid=grid, space_order=8) eq = Eq(u.forward, u.laplace + 1) # Ensure the Operator compiles correctly; there use to be a bug where # shared memory was attempted to be introduced in the gather/scatter # MPI routines Operator(eq, opt=opt4({'gpu-prefetch': 4, 'par-tile': (32, 4, 64)})) @switchconfig(platform='ampere') def test_noprefetch_w_raw_dep(self): grid = Grid(shape=(31, 17, 25)) so = (16, 8, 8) v = VectorTimeFunction(name='v', grid=grid, space_order=so) tau = TensorTimeFunction(name='tau', grid=grid, space_order=so) mu = Function(name='mu', grid=grid, space_order=so, avg_mode='safe_harmonic') v_x, v_y, _ = v _, tau_xy, _, _, _, _ = tau.values() eqn = Eq(tau_xy.forward, tau_xy + mu*(v_x.dy + v_y.dx)) opt = {'errctl': 'basic', 'gpu-fit': [tau_xy], 'par-tile': ((16, 8), (32, 8, 16)), 'gpu-vect': False, 'gpu-prefetch': 2} op0 = Operator(eqn, opt=optX({**opt, 'cire-hoist': False})) op1 = Operator(eqn, opt=optX(opt)) # Check generated code kernel0 = op1._func_table['kernel0'].root self._check_gencode(kernel0, 4, 2, 8) tau_xy1 = tau_xy.func(name='tau_xy1') mu.data[:] = 3.5 for i in (v_x, v_y, tau_xy, tau_xy1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, tau_xy=tau_xy1) diff = np.linalg.norm(np.asarray(tau_xy.data) - np.asarray(tau_xy1.data)) matched = np.allclose(tau_xy.data, tau_xy1.data, rtol=1.e-3) assert matched, "diff_norm=%e" % diff @switchconfig(platform='ampere') def test_etti_w_strain_like(self): grid = Grid(shape=(512, 512, 512)) x, y, z = grid.dimensions so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) e = TensorFunction(name="e", grid=grid, space_order=so, interp_order=16) v = VectorTimeFunction(name="v", grid=grid, space_order=so, time_order=1) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so, time_order=1) e_xx, e_xy, e_xz, e_yy, e_yz, e_zz = e.values() v_x, v_y, v_z = v.values() tau_xx, tau_xy, tau_xz, tau_yy, tau_yz, tau_zz = tau.values() eqns = [Eq(e_xy, v_x.forward + 1), Eq(e_xz, v_x.forward + 1), Eq(e_yz, v_y.forward + 1), Eq(e_zz, v_z.forward + 1), Eq(tau_xy.forward, e_xz + f*e_zz + tau_xy), Eq(tau_xz.forward, e_xy + e_zz + tau_xz), Eq(tau_yz.forward, e_xz + e_zz + tau_yz)] op = Operator(eqns, opt=optX({'errctl': 'basic', 'gpu-vect': False, 'cire-mode': ('automatic', 'automatic', 'x|y'), 'par-tile': (32, 8, 64)})) op.cfunction # Check generated code self._check_gencode(op._func_table['kernel0'].root, 4, 2, 8) self._check_gencode(op._func_table['kernel1'].root, 2, 1, 14) self._check_gencode(op._func_table['kernel2'].root, 4, 2, 30) @switchconfig(platform='ampere') def test_strain_interpolations_share_x_loop(self): grid = Grid(shape=(64, 64, 64)) x, _, _ = grid.dimensions so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) e = TensorFunction(name='e', grid=grid, space_order=so, interp_order=16) v = VectorTimeFunction(name='v', grid=grid, space_order=so, time_order=1) tau = TensorTimeFunction(name='tau', grid=grid, space_order=so, time_order=1) _, e_xy, e_xz, _, _, e_zz = e.values() v_x, _, v_z = v.values() _, tau_xy, tau_xz, _, tau_yz, _ = tau.values() eqns = [Eq(e_xy, v_x.forward + 1), Eq(e_zz, v_z.forward + 1), Eq(tau_xy.forward, e_xz + f*e_zz + tau_xy), Eq(tau_xz.forward, e_xy + e_zz + tau_xz), Eq(tau_yz.forward, e_xz + e_zz + tau_yz)] op0 = Operator(eqns, opt=opt0({'errctl': 'basic'})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'gpu-prefetch': 4, 'gpu-vect': False, 'cire-mode': ('automatic', 'automatic', 'x|y'), 'par-tile': (32, 8, 64)})) kernel = op1._func_table['kernel2'].root xloops = [i for i in FindNodes(Iteration).visit(kernel) if i.dim.root is x] shmas = {i for i in FindSymbols().visit(kernel) if i._mem_shared} assert len(shmas) == 2 assert len(xloops) == 1 exprs = FindNodes(Expression).visit(xloops[0]) reloads = [i for i in exprs if i.write._mem_registers and e_zz in i.reads] assert not reloads targets = (e_xy, e_zz, tau_xy, tau_xz, tau_yz) for i in (f, e_xz, v_x, v_z, *targets): init_chessboard(i) initial = tuple(np.array(i.data) for i in targets) op0.apply(time_M=3) expected = tuple(np.array(i.data) for i in targets) for i, data in zip(targets, initial, strict=True): i.data[:] = data op1.apply(time_M=3) for i, data in zip(targets, expected, strict=True): assert np.allclose(i.data, data, rtol=1.e-3) @switchconfig(platform='ampere') def test_prefetch_mixed_dtypes(self): grid = Grid(shape=(31, 17, 25)) x, _, _ = grid.dimensions so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so, dtype=np.uint16) u = TensorFunction(name='u', grid=grid, space_order=so, interp_order=2) v = TimeFunction(name='v', grid=grid, space_order=so, time_order=1) w = TensorTimeFunction(name='w', grid=grid, space_order=so, time_order=1) u_yz = u[1, 2] eqns = [Eq(u_yz, v.forward)] rhs = sin(f/65535)*u_yz eqns.extend(Eq(i.forward, rhs) for i in (w[0, 1], w[0, 2], w[1, 2])) op = Operator(eqns, opt=optX({'gpu-prefetch': 2})) op.cfunction kernel = op._func_table['kernel1'].root xloop, = [i for i in FindNodes(Iteration).visit(kernel) if i.dim.root is x] calls = [i for i in FindNodes(Call).visit(xloop) if i.name == '__pipeline_memcpy_async'] assert any('dcomm' in str(i) and 'u_yz' in str(i) for i in calls) @switchconfig(platform='ampere') def test_strain_explicit_interpolations_like(self): grid = Grid(shape=(31, 17, 25)) _, y, z = grid.dimensions so = (16, 8, 8) aux = TensorFunction(name='aux', grid=grid, space_order=so, interp_order=2) v = TimeFunction(name='v', grid=grid, space_order=so, time_order=1) p = TensorTimeFunction(name='p', grid=grid, space_order=so, time_order=1) aux_yy, aux_yz = aux.values()[3:5] p_yy, p_yz = p.values()[3:5] # These two interpolation directions require opposite sides of the same # shared-memory halo when lowered. kw = {'deriv_order': (0, 0), 'fd_order': (2, 2)} to_center = Derivative(aux_yz, y, z, x0={y: y, z: z}, **kw) to_yz = Derivative(aux_yy, y, z, x0={y: y + y.spacing/2, z: z + z.spacing/2}, **kw) eqns = [Eq(aux_yy, v.forward), Eq(aux_yz, 2*v.forward), Eq(p_yy.forward, aux_yy + 2*to_center + p_yy), Eq(p_yz.forward, 2*to_yz + 4*aux_yz + p_yz)] options = {'errctl': 'basic', 'gpu-opt': True, 'gpu-vect': 'stream-minimal', 'gpu-shm': 'max-reuse', 'par-tile': (32, 8, 48)} op0 = Operator(eqns, name='reference', opt=('advanced', {**options, 'gpu-prefetch': False})) op1 = Operator(eqns, name='prefetch', opt=('advanced', {**options, 'gpu-prefetch': 4, 'gpu-glb-mem': 'async-pipe'})) kernel = op1._func_table['kernel1'].root async_loads = [i for i in FindNodes(Call).visit(kernel) if i.name == '__pipeline_memcpy_async'] assert any('aux_yz' in str(i) and 'float4' in str(i) for i in async_loads) aux_yy1 = aux_yy.func(name='aux_yy1') aux_yz1 = aux_yz.func(name='aux_yz1') p_yy1 = p_yy.func(name='p_yy1') p_yz1 = p_yz.func(name='p_yz1') for i in (v, p_yy, p_yz, p_yy1, p_yz1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, aux_yy=aux_yy1, aux_yz=aux_yz1, p_yy=p_yy1, p_yz=p_yz1) for i, i1 in ((p_yy, p_yy1), (p_yz, p_yz1)): assert np.allclose(i.data, i1.data, rtol=1.e-5) @pytest.mark.parametrize('gpu_vect', [False, 'stream-only']) def test_prefetch_middle_subdomain(self, gpu_vect): grid = Grid(shape=(9, 11, 13)) u = TimeFunction(name='u', grid=grid, space_order=0) op = Operator(Eq(u.forward, u + 1, subdomain=grid.interior), name='middle_field_prefetch', opt=optX({ 'gpu-prefetch': 4, 'gpu-vect': gpu_vect, 'par-tile': (32, 4, 8) })) self._check_gencode(op._func_table['kernel0'].root, 2, 1, 2) op.apply(time_M=1) expected = np.zeros_like(u.data) expected[0, 1:-1, 1:-1, 1:-1] = 2 expected[1, 1:-1, 1:-1, 1:-1] = 1 np.testing.assert_array_equal(u.data, expected) class TestLaunchBounds: @pytest.mark.parametrize('opt', [True, [True], [True, False]]) def test_basic(self, opt): grid = Grid(shape=(16, 16, 16)) u = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u + 1) op = Operator(eq, opt=opt0({'par-tile': (4, 4), 'gpu-launch-bounds': opt})) # Check generated code assert '__launch_bounds__(16)' in str(op) op.apply(time_M=10) assert np.all(u.data[0] == 10) assert np.all(u.data[1] == 11) def test_multiple_kernels(self): grid = Grid(shape=(16, 16, 16)) u = TimeFunction(name='u', grid=grid, space_order=4) v = TimeFunction(name='v', grid=grid, space_order=4) eqns = [Eq(u.forward, u + 1), Eq(v.forward, u.forward.dy + 2)] mask0 = [True, False] mask1 = [False, True] mask2 = [True, True] opt0 = opt4({'par-tile': (4, 4, 4), 'gpu-launch-bounds': mask0}) opt1 = opt4({'par-tile': (4, 4, 4), 'gpu-launch-bounds': mask1}) opt2 = opt4({'par-tile': (4, 4, 4), 'gpu-launch-bounds': mask2}) op0 = Operator(eqns, opt=opt0) op1 = Operator(eqns, opt=opt1) op2 = Operator(eqns, opt=opt2) # Check generated code for op, (m0, m1) in zip([op0, op1, op2], [mask0, mask1, mask2]): kernel0 = op._func_table['kernel0'].root kernel1 = op._func_table['kernel1'].root if m0: assert '__launch_bounds__(64)' in str(kernel0) else: assert '__launch_bounds__' not in str(kernel0) if m1: assert '__launch_bounds__(16)' in str(kernel1) else: assert '__launch_bounds__' not in str(kernel1) @skipif('nocuda') class TestVectorLoads: def _check_gencode_stream_loads(self, kernel, fields, exp_stream_loads): dims = FindSymbols('dimensions').visit(kernel) assert len([d for d in dims if isinstance(d, ThreadDimension)]) == 6 exprs = FindNodes(Expression).visit(kernel) streaming_loads = [e for e in exprs if e.expr.find(Cast) and set(fields) & set(e.functions)] assert len(streaming_loads) == exp_stream_loads def _check_gencode_prefetch(self, efunc, nasyncs): calls = FindNodes(Call).visit(efunc) for call, count in [('__pipeline_memcpy_async', nasyncs)]: found = [i for i in calls if i.name == call] assert len(found) == count # Ensure the smart thread remap is used dest, _, _ = found[0].arguments assert any(isinstance(i, ThreadDimLayoutB) for i in dest.free_symbols) @switchconfig(ignore_unknowns=True) @switchconfig(platform='ampere') @pytest.mark.parametrize('par_tile', [ (16, 8, 4), (32, 16, 16), ]) @pytest.mark.parametrize('exprs,exp_stream_loads', [ # Simple cases (['u.dy + f'], 1), (['u.dy + u.dz + f'], 3), # Fallback, to prioritize reuse via registers (['u.laplace + f'], 4), (['u.laplace + f', 'v.laplace + f'], 8), # More complex cases (['u.dy.dx + f + g'], 2), (['u.dx + u.dy.dz + u.dy.dx + f + g'], 4), (['u.dz2 + u.dy.dx'], 6), (['u.dx2 + u.dz2 + u.dy2 + u.dy.dx + u.dx.dz'], 4), (['u.dy.dx + u.dy.dz + f*u.dx.dz'], 4), (['u.dz2 + u.dy.dx + u.dy.dz'], 2), (['u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz'], 4), (['u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz', 'v.dx2 + v.dy2 + v.dz2 + v.dy.dx + v.dy.dz + f*v.dx.dz'], 8), ]) def test_basic(self, par_tile, exprs, exp_stream_loads): grid = Grid(shape=(41, 27, 24)) x, _, _ = grid.dimensions nt = 5 so = 8 f = Function(name='f', grid=grid, space_order=so) g = Function(name='g', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so, save=nt, layers=NoLayers) v = u.func(name='v') u1 = u.func(name='u1') v1 = v.func(name='v1') fields = [u, v] eqns = [] for expr, field in zip(exprs, fields): eqns.append(Eq(field.forward, eval(expr) + 1.)) op0 = Operator(eqns, opt=opt4({'gpu-fit': fields, 'par-tile': par_tile})) op1 = Operator(eqns, opt=opt4({'gpu-fit': fields, 'par-tile': par_tile, 'gpu-vect': 'stream-only'})) # Check generated code kernel = op1._func_table['kernel0'].root self._check_gencode_stream_loads(kernel, fields, exp_stream_loads) g.data[:] = 0.31 f.data[:] = 0.75 for i in (u, u1, v, v1): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1, v=v1) for f0, f1 in [(u, u1), (v, v1)]: assert np.all(f0.data == f1.data) @switchconfig(platform='ampere') def test_smart_thread_decomposition(self): grid = Grid(shape=(42, 21, 33)) u = TimeFunction(name='u', grid=grid, space_order=8) eqn = Eq(u.forward, u.laplace + .01) op = Operator(eqn, opt=optX({'par-tile': (32, 8, 24), 'gpu-vect': 'stream-only', 'gpu-prefetch': 3})) op.cfunction # Check generated code assert '__pipeline_wait_prior(1);' in str(op) assert '__pipeline_memcpy_async' in str(op) self._check_gencode_prefetch(op._func_table['kernel0'].root, 12) @switchconfig(platform='ampere') @pytest.mark.parametrize('expr,expected,extra', [ ('u.dx + f*u.dy.dx', 6, {}), ('u.laplace', 8, {}), ('u.dz2 + u.dy.dx + f*u.dx.dz', 6, {}), ('u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz', 6, {}), ('u.laplace', 12, {'gpu-prefetch': 3}), ]) def test_with_prefetching(self, expr, expected, extra): grid = Grid(shape=(42, 21, 33)) nt = 5 so = 8 f = Function(name='f', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so, save=nt, layers=NoLayers) u1 = u.func(name='u1') eqn = Eq(u.forward, eval(expr)) op0 = Operator(eqn, opt=optX({'gpu-fit': [u], 'par-tile': (16, 8, 4), 'gpu-vect': False})) op1 = Operator(eqn, opt=optX({**{'errctl': 'basic', 'gpu-fit': [u], 'par-tile': (16, 8, 4), 'gpu-vect': 'stream-only'}, **extra})) # Check generated code assert '__pipeline_wait_prior(1);' in str(op1) assert '__pipeline_memcpy_async' in str(op1) self._check_gencode_prefetch(op1._func_table['kernel0'].root, expected) f.data[:] = 0.75 for i in (u, u1): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.all(u.data == u1.data) @switchconfig(platform='ampere') @pytest.mark.parametrize('exprs', [ ['vx.dx + vx.dz + vx.dy'], ['vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + C'], ['vx.dz', 'vz.dx'], ['vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + C', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + C'], ['vx.dy', 'vx.dz', 'vx.dy'], ['vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + C', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + C', 'vx.dx + f*vy.dx + vx.dy + g*vz.dx + vx.dz + C'], ['vx.dz', 'vx.dx + vx.dy', 'vz.dx', 'vy.dy'], ['vx.dz', 'vx.dx', '1', 'vx.dy'], ['vx.dz + vx.dy', 'vx.dx + vx.dz', 'vx.dx + vx.dz'], ['vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + C', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + C', 'vx.dx + f*vy.dx + vx.dy + g*vz.dx + vx.dz + C', 'vy.dy + f*vx.dy + vy.dz + vy.dx + g*vz.dy + C', 'vy.dy + f*vx.dy + vy.dx + g*vz.dy + vy.dz + C', 'vz.dz + f*vx.dz + g*vy.dz + vz.dy + vz.dx + C'], ]) @pytest.mark.parametrize('gpu_vect', [ (False, 'stream-only'), (True, 'stream') ]) @switchconfig(ignore_unknowns=True) def test_elastic_tti_like_incremental(self, exprs, gpu_vect): grid = Grid(shape=(42, 19, 27)) x, y, z = grid.dimensions nt = 5 f = Function(name='f', grid=grid, space_order=8) g = Function(name='g', grid=grid, space_order=8) v = VectorTimeFunction(name="v", grid=grid, space_order=8) tau = TensorTimeFunction(name="tau", grid=grid, space_order=8, save=nt, layers=NoLayers) vx, vy, vz = v tauxx, tauxy, tauxz, tauyy, tauyz, tauzz = F = tau.values() C = .001*(tauxx + tauxy + tauxz + tauyy + tauyz + tauzz) # noqa eqns = [] for expr, field in zip(exprs, F): eqns.append(Eq(field.forward, eval(expr) + 1.)) op0 = Operator(eqns, opt=optX({'errctl': 'basic', 'gpu-fit': F, 'par-tile': (16, 8), 'gpu-vect': gpu_vect[0]})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'gpu-fit': F, 'par-tile': (16, 8), 'gpu-vect': gpu_vect[1]})) f.data[:] = 19.0 g.data[:] = 18.0 vx.data[:] = 11.0 vy.data[:] = 12.0 vz.data[:] = 7.0 overrides = {i.name: i.func(name='%s1' % i.name) for i in F} for objs in (F, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) for i, j in zip(F, overrides.values()): if gpu_vect[0]: # See issue #499 assert np.allclose(i.data[:nt-1], j.data[:nt-1], rtol=1.e-3) else: assert np.allclose(i.data, j.data, rtol=1.e-3) @switchconfig(platform='ampere') def test_two_kernels(self): grid = Grid(shape=(42, 19, 27)) nt = 5 f = Function(name='f', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, space_order=8) v = TimeFunction(name='v', grid=grid, space_order=8) w = TimeFunction(name='w', grid=grid, space_order=8) eqns = [Eq(u.forward, u.dy + f*v.dy + f*w.dz + .01), Eq(v.forward, u.forward.dy + v + w + .02)] op0 = Operator(eqns, opt=optX({'errctl': 'basic', 'par-tile': (16, 8), 'gpu-vect': False})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'par-tile': (16, 8), 'gpu-vect': 'stream'})) f.data[:] = .03 u1 = u.func(name='u1') v1 = v.func(name='v1') op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1, v=v1) assert np.all(u.data == u1.data) assert np.all(v.data == v1.data) @switchconfig(platform='ampere') def test_fused_guards(self): grid = Grid(shape=(41, 27, 24)) x, _, _ = grid.dimensions nt = 5 # Since `so = 16` and the y-par_tile value is only 4, we expect the # compiler to exploit the layoutB to avoid generating useless guards # along y so = 16 par_tile = (8, 4, 16) f = Function(name='f', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so, save=nt, layers=NoLayers) eqn = Eq(u.forward, u.dx + u.dy.dz + u.dy.dx + f + 1.) op0 = Operator(eqn, opt=opt4({'gpu-fit': u, 'par-tile': par_tile})) op1 = Operator(eqn, opt=opt4({'gpu-fit': u, 'par-tile': par_tile, 'gpu-vect': 'stream-only'})) # Check generated code kernel = op1._func_table['kernel0'].root self._check_gencode_stream_loads(kernel, [u], 4) f.data[:] = 0.75 init_chessboard(u) u1 = u.func(name='u1') init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.all(u.data == u1.data) @switchconfig(platform='ampere') def test_fused_guards_v2(self): grid = Grid(shape=(42, 19, 27)) x, y, z = grid.dimensions nt = 5 v = VectorTimeFunction(name="v", grid=grid, space_order=8) tau = TensorTimeFunction(name="tau", grid=grid, space_order=8, save=nt, layers=NoLayers) vy = v[1] tauxx = tau[0] eqn = Eq(tauxx.forward, vy.dy + vy.dz + tauxx + 1.) op0 = Operator(eqn, opt=opt4({'errctl': 'basic', 'gpu-fit': tauxx, 'par-tile': (16, 8)})) op1 = Operator(eqn, opt=optX({'errctl': 'basic', 'gpu-fit': tauxx, 'par-tile': (16, 8), 'gpu-vect': 'stream'})) # Check generated code kernel = op1._func_table['kernel0'].root self._check_gencode_prefetch(kernel, 2) vy.data[:] = 11.0 tauxx.data[:] = 100 tauxx1 = tauxx.func(name='tauxx1') tauxx1.data[:] = 100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, tau_xx=tauxx1) assert np.allclose(tauxx.data, tauxx1.data, rtol=1.e-3) @switchconfig(platform='ampere') def test_etti_velocity_like(self): TestAsyncPipeline().run_etti_velocity_like( 2, {'gpu-vect': 'stream-minimal'}, mode=1, expected=(3, 1, 39, 2) ) def test_two_pass_v0(self): grid = Grid(shape=(40, 40, 40)) x, y, z = grid.dimensions c11 = Function(name='c11', grid=grid, space_order=8) c44 = c11.func(name='c44') txy = TimeFunction(name='txy', grid=grid, space_order=8, time_order=1, save=Buffer(1)) txz = txy.func(name='txz') kw = dict(deriv_order=(0, 0), fd_order=(2, 2)) expr0 = Derivative(c11, x, z, x0={x: x + x.spacing/2, z: z + z.spacing/2}, **kw) expr1 = Derivative(c44, x, y, x0={x: x + x.spacing/2, y: y + y.spacing/2}, **kw) eqns = [Eq(txz.forward, expr0 + txz), Eq(txy.forward, expr1 + txy)] op0 = Operator(eqns, opt=opt4()) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': (32, 4), 'cire-mode': '2pass', 'gpu-vect': 'stream-minimal', 'gpu-flip-derivs': True})) txy1 = txy.func(name='txy1') txz1 = txz.func(name='txz1') c11.data[:] = 3.5 c44.data[:] = 6.2 for i in (txy, txz, txy1, txz1): init_chessboard(i) op0(time_M=5) op1(time_M=5, txy=txy1, txz=txz1) assert np.allclose(txy.data, txy1.data, rtol=1.e-4) assert np.allclose(txz.data, txz1.data, rtol=1.e-4) def test_two_pass_asymmetric_halo(self): grid = Grid(shape=(40, 40, 40)) x, _, z = grid.dimensions f = Function(name='f', grid=grid, space_order=16) u = TimeFunction(name='u', grid=grid, space_order=16, save=Buffer(1)) expr = Derivative(f, x, z, x0={x: x + x.spacing/2, z: z + z.spacing/2}, deriv_order=(0, 0), fd_order=(16, 16)) eq = Eq(u.forward, expr + u) options = {'errctl': 'basic', 'par-tile': (32, 4), 'cire-mode': '2pass', 'gpu-flip-derivs': True} op0 = Operator(eq, opt=opt4({**options, 'gpu-vect': False})) op1 = Operator(eq, opt=opt4({**options, 'gpu-vect': 'stream-minimal'})) # The 7-point left halo, 32-point core, and 8-point right halo are # loaded as one 48-point vectorized region kernel = op1._func_table['kernel1'].root shma, = [i for i in FindSymbols().visit(kernel) if isinstance(i, TBShm)] loads = [i for i in FindNodes(Expression).visit(kernel) if i.write is shma] assert len(loads) == 1 assert 'Tz4 <= 44' in str(kernel) f.data[:] = 1.5 u1 = u.func(name='u1') for i in (u, u1): init_chessboard(i) op0(time_M=1) op1(time_M=1, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) def test_eiso_like(self): grid = Grid(shape=(40, 40, 40)) x, y, z = grid.dimensions so = (16, 8, 8) v = VectorTimeFunction(name="v", grid=grid, space_order=so) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so) mu = Function(name='mu', grid=grid, space_order=so, avg_mode='safe_harmonic') expr_tau = tau.dt - diag(div(v)) - mu eq = Eq(tau.forward, solve(expr_tau, tau.forward)) extra = {'errctl': 'basic', 'gpu-vect': 'stream-minimal', 'gpu-flip-derivs': True, 'gpu-prefetch': (True, False)} # The `(31,4)` par-tile for the CIRE-generated kernel0 will cause the # allocation of extra loader threads to avoid misalignment op0 = Operator(eq, opt=opt4({'par-tile': ((31, 4), (16, 8)), **extra})) op0.apply(dt=0.001, time_M=1) # Ensure it compiles and runs without errors op1 = Operator(eq, opt=opt4({'par-tile': ((32, 4), (16, 8)), **extra})) op1.apply(dt=0.001, time_M=1) def test_temps_alignment_upon_2pass(self): grid = Grid(shape=(35, 45, 55)) space_order = 8 f = Function(name='f', grid=grid, space_order=space_order) u = TimeFunction(name='u', grid=grid, space_order=space_order) eqn = Eq(u.forward, cos(f).dz + u.dz.dz + f*u.dy.dz + 1) op = Operator(eqn, opt=opt4({'par-tile': (24, 8, 16), 'gpu-vect': 'stream-minimal'})) # Check generated code # NOTE: we're ensuring the base pointer is shifted by the expected amount alloc_on_device = op._func_table['alloc_on_device0'].root assert 'r0 += 24' in str(alloc_on_device) assert 'r0_vec - 24' in str(op.body.frees[1]) op.apply(time_M=0) @pytest.mark.parametrize('space_order', range(1, 33)) def test_alignment(self, space_order): """ Test vector alignment over a range of space orders """ grid = Grid(shape=(29, 29, 29)) x, y, z = grid.dimensions t = grid.time_dim v_x = TimeFunction( name='v_x', grid=grid, staggered=(0, 1, 0, 0), dimensions=(t, x, y, z), space_order=space_order ) v_y = v_x._rebuild(name='v_y', staggered=(0, 0, 1, 0)) v_z = v_x._rebuild(name='v_z', staggered=(0, 0, 0, 1)) p_xx = TimeFunction( name='p_xx', grid=grid, staggered=(0, 0, 0, 0), dimensions=(t, x, y, z), space_order=space_order ) p_xy = p_xx._rebuild(name='p_xy', staggered=(0, 1, 1, 0)) p_xz = p_xx._rebuild(name='p_xz', staggered=(0, 1, 0, 1)) p_yy = p_xx._rebuild(name='p_yy') p_yz = p_xx._rebuild(name='p_yz', staggered=(0, 0, 1, 1)) p_zz = p_xx._rebuild(name='p_zz') eqns = [ Eq(v_x, p_xx.dx + p_xz.dz), Eq(v_y, p_xy.dx + p_yy.dy + p_yz.dz), Eq(v_z, p_xz.dx + p_yz.dy + p_zz.dz) ] op = Operator(eqns, opt=('advanced', { 'errctl': 'basic', 'gpu-opt': True, 'gpu-vect': 'stream-minimal', 'par-tile': (64, 1) })) _ = op(dt=1e-5) # These options require Nvidia Ampere and newer @switchconfig(platform='ampere') @pytest.mark.parametrize('gpu_glb_mem', [ False, pytest.param('async-pipe', marks=skipif('nocuda-cc80')) ]) @pytest.mark.parametrize('gpu_prefetch', [ False, pytest.param(2, marks=skipif('nocuda-cc80')) ]) @pytest.mark.parametrize('dtype', [np.float16, np.complex64]) @pytest.mark.parametrize('so', [8, 16]) @pytest.mark.parametrize('par_tile', [(32, 8, 32), (16, 8, 16), (32, 8)]) def test_heterogeneous_vector_types(self, monkeypatch, gpu_glb_mem, gpu_prefetch, dtype, so, par_tile): grid = Grid(shape=(101, 101, 101), extent=(2*np.pi, 2*np.pi, 2*np.pi)) b = Function(name='b', grid=grid, space_order=so) txx = Function(name='txx', grid=grid, space_order=so, dtype=dtype) tyy = Function(name='tyy', grid=grid, space_order=so, dtype=dtype) tzz = Function(name='tzz', grid=grid, space_order=so, dtype=dtype) vx = Function(name='vx', grid=grid, space_order=so) # Initialise fields b.data[:] = 1 sinusoid = np.sin(2*np.pi*np.linspace(-so, 100+so, 101+(2*so))/100) txx.data_with_halo[:] = sinusoid[:, np.newaxis, np.newaxis] tyy.data_with_halo[:] = sinusoid[np.newaxis, :, np.newaxis] tzz.data_with_halo[:] = sinusoid[np.newaxis, np.newaxis, :] rhs = b*(txx.dx + tyy.dy + tzz.dz) if dtype == np.complex64: rhs = Real(rhs) eq = Eq(vx, rhs) warnings = [] monkeypatch.setattr(cuda, 'event_suboptimal_kernel', warnings.append) # b and txx, tyy, tzz have different dtypes which will lead to # heterogeneous vector types: `float16` -> `__half4` (actually a # `float2`), `complex64` -> `double4`, and `float32` -> `float4`. # Note that some combinations of optimisations will result in only # half-precision vector loads being created. op = Operator(eq, opt=optX({'errctl': 'basic', 'par-tile': par_tile, 'gpu-vect': 'stream-minimal', 'gpu-prefetch': gpu_prefetch, 'gpu-glb-mem': gpu_glb_mem})) # Whilst the "official" longest vector type for FP16 is `__half2`, we # load 4 halves as a single `float2` - in effect a `__half4 z_size = par_tile[0] y_size = par_tile[1] assert f"int Tz4 = 4*(((threadIdx.x) & {z_size//2 - 1}));" in str(op.ccode) # Vector load type. vname = 'float2' if dtype == np.float16 else 'double4' if gpu_glb_mem and gpu_prefetch: x = "x0_blk0" if len(par_tile) == 3 else "x_m" # tyy gets vector loaded aload = (f"__pipeline_memcpy_async(&s_yz0[({x} & 1)][Ty][Tz4]," f"&tyyL0({x} + {so}, Gy + {so//2}, Gz4 + {so})," f"sizeof({vname}));") sload = (f"*(reinterpret_cast<{vname}*>(&s_yz0[({x} & 1)][Ty][Tz4])) = " f"*(reinterpret_cast<{vname}*>(&tyyL0({x} + {so}, " f"Gy + {so//2}, Gz4 + {so})));") if dtype == np.float16: assert aload in str(op.ccode) else: assert aload not in str(op.ccode) assert sload in str(op.ccode) assert (f"__shared__ float s_yz2[2][{y_size}][{z_size}] __attribute__ " "((aligned (64)));" in str(op.ccode)) assert ("__pipeline_memcpy_async(&s_yz2[dcomm][Ty][Tz4]," f"&bL0(x + {so+1}, Gy + {so}, Gz4 + {so}),sizeof(float4));" in str(op.ccode)) if gpu_glb_mem and not gpu_prefetch: aload = ("__pipeline_memcpy_async(&s_yz1[Ty][Tz4]," f"&tzzL0(x + {so}, Gy + {so}, Gz4 + {so//2})," f"sizeof({vname}));") sload = (f"*(reinterpret_cast<{vname}*>(&s_yz1[Ty][Tz4])) = " f"*(reinterpret_cast<{vname}*>(&tzzL0(x + {so}, " f"Gy + {so}, Gz4 + {so//2})));") if dtype == np.float16: assert aload in str(op.ccode) else: assert aload not in str(op.ccode) assert sload in str(op.ccode) fallback_warning = ( "asynchronous pipeline supports at most 16-byte copies; " "using synchronous loads for 32-byte vector types" ) if gpu_glb_mem and dtype == np.complex64: assert warnings == [fallback_warning] else: assert not warnings # Only half-precision vector loads without async pipelines if not gpu_glb_mem and not gpu_prefetch: assert (f"*(reinterpret_cast<{vname}*>(&s_yz0[Ty][Tz4])) = " f"*(reinterpret_cast<{vname}*>(" f"&tyyL0(x + {so}, Gy + {so//2}, Gz4 + {so})));" in str(op.ccode)) op.apply() # Check operator runs cosinusoid = np.cos(2*np.pi*np.linspace(0, 100, 101)/100) check = cosinusoid[:, np.newaxis, np.newaxis] \ + cosinusoid[np.newaxis, :, np.newaxis] \ + cosinusoid[np.newaxis, np.newaxis, :] assert np.all(np.isclose(vx.data, check, atol=2e-2)) @switchconfig(platform='ampere') @pytest.mark.parametrize('expr', ['Eq(f, f*g)', 'Eq(f, f*g.dx)', 'Eq(f, f*g.dy)', 'Eq(f, f*g.dz)']) @pytest.mark.parametrize('gpu_vect', ['stream-minimal', False]) @pytest.mark.parametrize('dtype', [Float8(0., 4.5), np.float16]) @pytest.mark.parametrize('so', [4, 8]) def test_issue_809(self, expr, dtype, gpu_vect, so): grid = Grid(shape=(11, 11, 11)) x, y, _ = grid.dimensions f = TimeFunction(name='f', grid=grid, space_order=so, # noqa staggered=(x, y), save=Buffer(1)) # noqa g = Function(name='g', grid=grid, space_order=so, # noqa staggered=NODE, dtype=dtype) # noqa op = Operator(eval(expr), opt=optX({'errctl': 'basic', 'gpu-vect': gpu_vect, 'gpu-prefetch': (True, False), 'par-tile': ((32, 8), (16, 8))})) # A single shared memory array should be generated for this setup. A second # shared memory array is generated where g.dz is present expected_shm = 2 if "dz" in expr else 1 assert len([f for f in FindSymbols().visit(op) if isinstance(f, TBShm)]) == expected_shm # This should have the correct dtype dtypename = "__half" if dtype == np.float16 else "unsigned char" shm_no = 1 if "dz" in expr else 0 assert f"__shared__ {dtypename} s_yz{shm_no}" in str(op) # Additionally, a single queue will be generated which will store the # promoted (and scaled in the Float8 case) function multiplied by # the interpolation weights. assert len([f for f in FindSymbols().visit(op) if isinstance(f, TBArray)]) == expected_shm + 1 if dtype != np.float16 or not gpu_vect: assert "__pipeline_memcpy_async" not in str(op) else: # FP16 should use async pipelines with float2 (fake __half4) # vector loads if available. assert "__pipeline_memcpy_async" in str(op) # Belt and braces - check that an erroneous reinterpret cast # on the reduced-precision array access is not generated as it # can produce a CUDA runtime error. if gpu_vect: assert "reinterpret_cast" not in str(op) op.apply(dt=0.1, time_M=100) @pytest.mark.parametrize('tkn', [4, 5, 13]) def test_subdomains_v0(self, tkn): class SD(SubDomain): name = 'inner' def define(self, dimensions): return {d: ('middle', tkn, tkn) for d in dimensions} grid = Grid(shape=(100, 100, 100)) sd = SD(grid=grid) u = TimeFunction(name='u', grid=grid, space_order=8) u1 = u.func(name='u1') eq = Eq(u.forward, u.dy + u.dz + 1., subdomain=sd) op0 = Operator(eq, opt=opt4({'par-tile': (16, 4, 4), 'gpu-vect': False})) op1 = Operator(eq, opt=optX({'errctl': 'basic', 'par-tile': (16, 4, 4), 'gpu-vect': 'stream-minimal'})) # Check generated code -- ensure we're using vector loads, regardless # of the actual `tkn` value assert ('__pipeline_memcpy_async(&s_yz0[dcomm][Ty][Tz4 + 4],' '&uL0(t0, ix + 9, Gy + 4, Gz4 + 8),' 'sizeof(float4))') in str(op1) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=3) op1.apply(time_M=3, u=u1) n0 = np.linalg.norm(u.data) n1 = np.linalg.norm(u1.data) assert np.isfinite(n0) assert np.isfinite(n1) assert np.isclose(n0, n1, rtol=1.e-6) @pytest.mark.parametrize('storage_margin', [0, 12]) @pytest.mark.parametrize('side', ['left', 'right', 'middle']) @pytest.mark.parametrize('axis', [1, 2]) @pytest.mark.parametrize('par_tile', [(32, 4, 48), (32, 16, 48)]) @pytest.mark.parametrize('prefetch,gpu_vect', [ (False, False), (2, 'stream-only'), (4, 'stream-only') ]) def test_subdomain_local_fields(self, storage_margin, side, axis, par_tile, prefetch, gpu_vect, host_kwargs): grid = Grid(shape=(32, 37, 50), extent=(31, 36, 49)) spec = ['none'] * 3 spec[axis] = side sd = SpeccedDomain('boundary', spec, 10, grid=grid) thickness = (max(0, 10 - storage_margin) if side == 'middle' else 10 + storage_margin) storage_sd = (SpeccedDomain('storage', spec, thickness, grid=grid) if storage_margin else sd) work = Function(name='work', grid=storage_sd, space_order=8) u = TimeFunction(name='u', grid=grid, space_order=8) u1 = u.func(name='u1') eq = Eq(u.forward, u + .01*(work.dx + work.dy + work.dz), subdomain=sd) op0 = Operator(eq, name='subdomain_local_reference', opt='noop', **host_kwargs) op1 = Operator(eq, name='subdomain_local_streaming', opt=('advanced', { 'gpu-opt': True, 'gpu-vect': gpu_vect, 'gpu-glb-mem': 'async-pipe', 'gpu-prefetch': prefetch, 'gpu-wdr': False, 'gpu-hoist': 2, 'par-tile': par_tile, 'fission': False })) if axis == 1 and gpu_vect: # The missing Y guard must be restored without losing vector LOADs assert any('__pipeline_memcpy_async' in line and 'workL0' in line and 'sizeof(float4)' in line for line in str(op1).splitlines()) rng = np.random.default_rng(0) for shape in (None, (33, 29, 41), (35, 45, 58)): inputs = {} if shape is not None: # Reuse both Operators with smaller and larger runtime fields runtime_grid = Grid(shape=shape, extent=tuple(n - 1 for n in shape)) runtime_sd = SpeccedDomain('storage', spec, thickness, grid=runtime_grid) work = Function(name='runtime_work', grid=runtime_sd, space_order=8) u = TimeFunction(name='runtime_u', grid=runtime_grid, space_order=8) u1 = u.func(name='runtime_u1') inputs = {'work': work, 'u': u} work.data_with_halo[:] = rng.random(work.data_with_halo.shape) u.data[:] = 1 u1.data[:] = 1 op0.apply(time_M=2, **inputs) op1.apply(time_M=2, **(inputs | {'u': u1})) np.testing.assert_allclose(u.data, u1.data, rtol=1.e-6, atol=1.e-6) @pytest.mark.parametrize('par_tile', [(32, 4), (32, 4, 48)]) @pytest.mark.parametrize('subdomain', [None, 'interior']) def test_2pass_v0(self, par_tile, subdomain): grid = Grid(shape=(100, 37, 50)) subdomain = grid.subdomains.get(subdomain) so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so) v = u.func(name='v') eqns = [ Eq(u.forward, u.dy.dx + u.dy.dz + f*u.dx.dz, subdomain=subdomain), Eq(v.forward, v.dy.dx + v.dy.dz + f*v.dx.dz, subdomain=subdomain) ] op0 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': par_tile})) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': par_tile, 'cire-mode': '2pass', 'gpu-vect': 'stream-minimal'})) u1 = u.func(name='u1') v1 = v.func(name='v1') for i in (u, v, u1, v1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, u=u1, v=v1) assert np.allclose(u.data, u1.data, rtol=1.e-4) assert np.allclose(v.data, v1.data, rtol=1.e-4) def test_2pass_v1(self): grid = Grid(shape=(100, 37, 50)) so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so) v = u.func(name='v') eqns = [ Eq(u.forward, u.dy.dx + u.dy.dz + f*u.dx.dz), Eq(v.forward, v.dy.dx + v.dy.dz + f*v.dx.dz) ] op0 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': (32, 4, 48)})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'par-tile': (32, 4, 48), 'cire-mode': 'x|y', 'gpu-prefetch': 4, 'gpu-vect': 'stream-minimal'})) u1 = u.func(name='u1') v1 = v.func(name='v1') for i in (u, v, u1, v1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, u=u1, v=v1) assert np.allclose(u.data, u1.data, rtol=1.e-2) assert np.allclose(v.data, v1.data, rtol=1.e-2) @switchconfig(platform='ampere') @pytest.mark.parametrize('nz', [24, 29]) def test_padded_snapshot_buffer(self, nz): """ A snapshot is auto-padded along the innermost Dimension, so that the vector LOADs reading it back are aligned whatever the grid size. """ grid = Grid(shape=(48, 48, nz)) nt = 4 so = 8 usave = TimeFunction(name='usave', grid=grid, space_order=so, save=nt) v = TimeFunction(name='v', grid=grid, space_order=so) g = Function(name='g', grid=grid, space_order=so) eqns = [Inc(g, v*usave.laplace), Eq(v.backward, v.laplace + 1)] op = Operator(eqns, opt=optX({'errctl': 'basic', 'par-tile': (32, 4, 32), 'gpu-vect': 'stream-minimal'})) # Check generated code -- every buffer must carry the padded extent padded = 'ROUND_UP(z_size + 16, 32)' assert padded in str(op) assert 'z_size + 16)*sizeof(float)' not in str(op) # And the LOADs off the device buffer are indeed vectorized assert any('usaveb0' in i and 'float4' in i for i in str(op).split('\n')) @switchconfig(platform='ampere') @switchconfig(autopadding=False) def test_unpadded_snapshot_buffer(self): """ Without auto-padding nothing keeps a snapshot's rows aligned, so no vector LOAD may be synthesized off it. """ grid = Grid(shape=(48, 48, 29)) nt = 4 so = 8 usave = TimeFunction(name='usave', grid=grid, space_order=so, save=nt) v = TimeFunction(name='v', grid=grid, space_order=so) g = Function(name='g', grid=grid, space_order=so) eqns = [Inc(g, v*usave.laplace), Eq(v.backward, v.laplace + 1)] op = Operator(eqns, opt=optX({'errctl': 'basic', 'par-tile': (32, 4, 32), 'gpu-vect': 'stream-minimal'})) # Check generated code -- no vector type alongside `usaveb0` for line in str(op).split('\n'): if 'usaveb0' in line: assert not any(i in line for i in ['float2', 'float3', 'float4']) @switchconfig(platform='ampere') def test_streaming_guarded_core(self): grid = Grid(shape=(31, 17, 65)) _, y, z = grid.dimensions so = (16, 8, 8) f = Function(name='f', grid=grid, space_order=so) q = Function(name='q', grid=grid, space_order=so, dtype=np.uint8) u = TimeFunction(name='u', grid=grid, space_order=so) u1 = u.func(name='u1') left = Derivative(f, z, x0=z - z.spacing/2, deriv_order=0, fd_order=2) cross = Derivative(q*u, y, z, x0={y: y + y.spacing/2, z: z + z.spacing/2}, deriv_order=(0, 0), fd_order=(2, 2)) eqn = Eq(u.forward, u + left + cross) # Right loader threads serve the cross interpolation. The separate # left halo of `f` must be fused with only its 32 active core threads options = {'errctl': 'basic', 'gpu-opt': True, 'gpu-shm': 'min-reuse', 'gpu-loader-threads-radius': (0, 4), 'cire-hoist': False, 'par-tile': (32, 8, 48)} op0 = Operator(eqn, name='reference', opt=('advanced', {**options, 'gpu-vect': False, 'gpu-prefetch': False})) op1 = Operator(eqn, name='prefetch', opt=('advanced', {**options, 'gpu-vect': 'stream-minimal', 'gpu-glb-mem': 'async-pipe', 'gpu-prefetch': 4})) assert sum('Ty <= 7' in i and 'Tz4 <= 32' in i for i in str(op1).splitlines()) == 2 assert all(i not in str(op1) for i in ('Ty >= Ty', 'Tz4 >= Tz4')) for i in (f, u, u1): init_chessboard(i) q.data[:] = 1 op0.apply(time_M=3) op1.apply(time_M=3, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) class TestCooperativeReductions: @pytest.mark.parametrize('beta,options,expected', [ (0, {}, False), (1, {}, True), (2, {}, True), (3, {}, True), (0, {'gpu-opt': True}, True), (1, {'gpu-opt': False}, False), (0, {'gpu-reduce-local': True}, True), (1, {'gpu-reduce-local': False}, False), ]) def test_default(self, beta, options, expected): grid = Grid(shape=(17, 19, 21)) f = Function(name='f', grid=grid, space_order=8) rcv = SparseFunction(name='rcv', grid=grid, npoint=65, interpolation='sinc', r=4) with switchconfig(beta=beta): op = Operator(rcv.interpolate(expr=f), opt=('advanced', options)) assert ('__shfl_down' in str(op)) is expected @pytest.mark.parametrize('r', [3, 4], ids=['padded', 'full']) @pytest.mark.parametrize('dtype', [np.float32, np.float64]) def test_interpolate(self, r, dtype, host_kwargs): grid = Grid(shape=(33, 35, 37), extent=(32., 34., 36.), dtype=dtype) p = TimeFunction(name='p', grid=grid, space_order=8, time_order=2) q = p.func(name='q') rcv = SparseTimeFunction(name='rcv', grid=grid, npoint=65, nt=3, interpolation='sinc', r=r) eqns = rcv.interpolate(expr=p.forward/3 + 2*q.forward/3, increment=True) op0 = Operator(eqns, name='CooperativeHost', **host_kwargs) op1 = Operator(eqns, name='Cooperative', opt=opt0({'gpu-reduce-local': True})) assert '__shfl_down' in str(op1) init_chessboard(p) q.data_with_halo[:] = 1/p.data_with_halo rcv.coordinates.data[:] = np.linspace(0, grid.extent, rcv.npoint) rcv.coordinates.data[0] = 0 rcv.coordinates.data[1] = grid.extent # A valid receiver in the last, partially populated block rcv.coordinates.data[-1] = np.asarray(grid.extent)/2 bounds = dict(x_m=12, x_M=20, y_m=12, y_M=22, z_m=12, z_M=24, time_M=2) rcv.data[:] = 2.5 op0.apply(**bounds) expected = rcv.data.copy() rcv.data[:] = 2.5 op1.apply(**bounds) assert np.all(expected[:, :2] == 2.5) assert np.any(expected != 2.5) tol = 1.e-6 if dtype == np.float32 else 1.e-13 assert np.allclose(rcv.data, expected, rtol=10*tol, atol=tol) @pytest.mark.parametrize('variant', ['abs', 'modulo']) def test_wrapped_index(self, variant, host_kwargs): grid = Grid(shape=(33,), extent=(32.,)) x, = grid.dimensions f = Function(name='f', grid=grid, space_order=8) rcv = SparseFunction(name='rcv', grid=grid, npoint=4, interpolation='sinc', r=4) index = sympy.Abs(x) if variant == 'abs' else sympy.Mod(x, 16) eqns = rcv.interpolate(f._subs(x, index)) op0 = Operator(eqns, name='WrappedHost', **host_kwargs) op1 = Operator(eqns, name='Wrapped', opt=opt0({'gpu-reduce-local': True})) assert '__shfl_down' in str(op1) f.data[:] = np.arange(33) rcv.coordinates.data[:, 0] = [0., .25, 1.5, 8.25] op0.apply() expected = rcv.data.copy() rcv.data[:] = np.nan op1.apply() assert np.allclose(rcv.data, expected, rtol=1.e-5, atol=1.e-6) @pytest.mark.parametrize('interpolation,r', [('linear', 1), ('sinc', 4)]) def test_one_dimension(self, interpolation, r, host_kwargs): grid = Grid(shape=(97,), extent=(96.,)) f = Function(name='f', grid=grid, space_order=8) rcv = SparseFunction(name='rcv', grid=grid, npoint=65, interpolation=interpolation, r=r) eqns = rcv.interpolate(expr=f) op0 = Operator(eqns, name='LocalSumHost', **host_kwargs) op1 = Operator(eqns, name='LocalSum', opt=opt0({'gpu-reduce-local': True})) # Kernel construction must keep both point and thread axes parallel assert len(FindNodes(KernelLaunch).visit(op1)[0].block.cargs) == 2 f.data_with_halo[:] = (np.arange(f.data_with_halo.size) % 7 - 3)/7. rcv.coordinates.data[:, 0] = np.linspace(0., 95.75, rcv.npoint) op0.apply(x_m=12, x_M=84) expected = rcv.data.copy() rcv.data[:] = np.nan op1.apply(x_m=12, x_M=84) assert np.allclose(rcv.data, expected, rtol=1.e-5, atol=1.e-6) @pytest.mark.parametrize('variant', ['disabled', 'strided', 'inject', 'dependent', 'nonzero']) @switchconfig(beta=0) def test_fallback(self, variant): grid = Grid(shape=(17, 19, 21)) x, y, z = grid.dimensions dimensions = (x, z, y) if variant == 'strided' else (x, y, z) shape = tuple(grid.shape[d] for d in dimensions) f = Function(name='f', grid=grid, dimensions=dimensions, shape=shape, space_order=8) rcv = SparseFunction(name='rcv', grid=grid, npoint=65, interpolation='sinc', r=4) if variant == 'inject': eqns = rcv.inject(field=f, expr=rcv) else: eqns = rcv.interpolate(expr=f).evaluate if variant in ('dependent', 'nonzero'): # These recurrences are deliberately outside LocalSum semantics result = eqns[-1] value = Symbol(name='sum', dtype=f.dtype) rhs = result.rhs.expr initial = 2.5 if variant == 'nonzero' else 0 if variant == 'dependent': rhs += value eqns = [*eqns[:-1], Eq(value, initial, implicit_dims=result.implicit_dims), Inc(value, rhs, implicit_dims=result.implicit_dims), result.func(result.lhs, value, implicit_dims=result.implicit_dims)] op = Operator(eqns, opt=opt0({'gpu-reduce-local': variant != 'disabled'})) assert '__shfl_down' not in str(op) @pytest.mark.parametrize('storage', ['scalar', 'tensor']) def test_setup(self, storage, host_kwargs): grid = Grid(shape=(97,), extent=(96.,)) f = Function(name='f', grid=grid) rcv = SparseFunction(name='rcv', grid=grid, npoint=65) pd = rcv._sparse_dim if storage == 'scalar': scale = Symbol(name='scale', dtype=f.dtype) else: scale = Function(name='scale', dimensions=(pd,), shape=(rcv.npoint,)) eqns = [Eq(scale, 2 + rcv, implicit_dims=(pd,)), rcv.interpolate(f*scale)] op0 = Operator(rcv.interpolate(f), name='SetupHost', **host_kwargs) op1 = Operator(eqns, name='Setup', opt=opt0({'gpu-reduce-local': True})) # Only private scalar setup can be replicated across threads assert ('__shfl_down' in str(op1)) is (storage == 'scalar') f.data_with_halo[:] = (np.arange(f.data_with_halo.size) % 7 - 3)/7. rcv.coordinates.data[:, 0] = np.linspace(0., 95.75, rcv.npoint) initial = np.linspace(0.5, 2.5, rcv.npoint) op0.apply() expected = (2 + initial)*rcv.data.copy() rcv.data[:] = initial op1.apply() assert np.allclose(rcv.data, expected, rtol=1.e-5, atol=1.e-6) @pytest.mark.parametrize('size,shuffles', [(6, 3), (64, 6)]) @pytest.mark.parametrize('dtype', [np.float32, np.float64]) @switchconfig(platform='amdgpuX', language='hip') def test_codegen_hip(self, size, shuffles, dtype): grid = Grid(shape=(97,), dtype=dtype) f = Function(name='f', grid=grid, space_order=size//2) # Precomputed weights allow a full AMD wavefront, beyond the sinc limit rcv = PrecomputedSparseFunction( name='rcv', grid=grid, npoint=65, r=size, gridpoints=np.zeros((65, 1), dtype=np.int32), interpolation_coeffs=np.full((65, 1, size), 1/size, dtype=dtype) ) op = Operator(rcv.interpolate(f), opt=opt0({'gpu-reduce-local': True})) assert str(op).count('__shfl_down(') == shuffles @switchconfig(platform='nvidiaX', language='sycl') def test_codegen_sycl(self): grid = Grid(shape=(97,)) f = Function(name='f', grid=grid) rcv = SparseFunction(name='rcv', grid=grid, npoint=65) op = Operator(rcv.interpolate(f), opt=opt0({'gpu-reduce-local': True})) assert all(i not in str(op) for i in ('SubgroupReduce', '__shfl_down')) @skipif('nocuda') class TestMaximizeILP: @switchconfig(platform='ampere') @pytest.mark.parametrize('expr', [ 'u.dx2 + u.dy2 + u.dz2 + u.dy.dz + u.dy.dx + u.dy.dx', ]) @pytest.mark.parametrize('gpu_prefetch', [False, 4]) def test_basic(self, expr, gpu_prefetch): grid = Grid(shape=(63, 42, 35)) nt = 5 u = TimeFunction(name='u', grid=grid, space_order=16) eqn = Eq(u.forward, eval(expr) + .01) op0 = Operator(eqn, opt=opt4({'par-tile': (32, 4, 48)})) op1 = Operator(eqn, opt=opt4({'par-tile': (32, 4, 48), 'gpu-prefetch': gpu_prefetch, 'gpu-max-ilp': 2})) # Check generated code assert str(op0) != str(op1) u1 = u.func(name='u1') init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-3) @switchconfig(platform='ampere') @pytest.mark.parametrize('gpu_prefetch', [False, 4]) def test_advanced(self, gpu_prefetch): grid = Grid(shape=(63, 42, 35)) x, y, z = grid.dimensions nt = 5 vx = TimeFunction(name='vx', grid=grid, space_order=8, time_order=1, save=Buffer(1)) txx = vx.func(name='txx') txy = vx.func(name='txy') txz = vx.func(name='txz') x0m = x - x.spacing/2 y0m = y - y.spacing/2 z0m = z - z.spacing/2 x0p = x + x.spacing/2 y0p = y + y.spacing/2 z0p = z + z.spacing/2 kw = dict(deriv_order=(0, 0), fd_order=(2, 2)) expr0 = (Derivative(txx.dx(x0=x0m), x, y, x0={x: x0p, y: y0m}, **kw) + txx.dy(x0=y0m) + Derivative(txx.dz(x0=z0m), y, z, x0={y: y0m, z: z0p}, **kw)) expr1 = (Derivative(txy.dx(x0=x0m), x, z, x0={x: x0p, z: z0m}, **kw) + Derivative(txy.dy(x0=y0m), y, z, x0={y: y0p, z: z0m}, **kw) + txy.dz(x0=z0m)) expr2 = (txz.dx(x0=x0p) + Derivative(txz.dy(x0=y0p), x, y, x0={x: x0p, y: y0m}, **kw) + Derivative(txz.dz(x0=z0p), x, z, x0={x: x0p, z: z0m}, **kw)) eqn = Eq(vx.forward, expr0 + expr1 + expr2 + vx) op0 = Operator(eqn, opt=opt4({'par-tile': (32, 4, 48), 'deriv-unroll': 'inner'})) op1 = Operator(eqn, opt=opt4({'par-tile': (32, 4, 48), 'gpu-prefetch': gpu_prefetch, 'deriv-unroll': 'inner', 'gpu-max-ilp': 2})) # Check generated code assert str(op0) != str(op1) vx1 = vx.func(name='vx1') for i in (vx, vx1, txx, txy, txz): init_chessboard(i) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, vx=vx1) assert np.allclose(vx.data, vx1.data, rtol=1.e-4) @pytest.mark.parametrize('sizes,access,unroll', [ ((8, 8, 8), 'write', 4), ((8, 8, 8), 'read', 4), ((8, 8, 8), 'dense', 0), ((4, 4, 4), 'write', 4), ((6, 6, 6), 'write', 4), ((8,), 'write', 0), ((8, 8), 'write', 0), ((8, 8, 8, 8), 'write', 4), ((8, 8, 8, 8, 8), 'write', 4), ((16, 8, 8), 'write', 4), ((8, 16, 8), 'write', 4), ((8, 8, 16), 'write', 4), ((16, 16, 16), 'write', 4), ((17, 8, 8), 'write', 0), ((8, 17, 8), 'write', 0), ((8, 8, 17), 'write', 0), ((None, 8, 8), 'write', 0), ((8, None, 8), 'write', 0), ((8, 8, None), 'write', 0), ((8, 8, 8), 'pragma', 1), ]) def test_unroll_many_short(self, sizes, access, unroll): grid = Grid(shape=(9, 9, 9)) rcv = SparseFunction(name='rcv', grid=grid, npoint=3) f = Function(name='f', grid=grid) n = Constant(name='n', value=8, dtype=np.int32) lhs = rcv[0] if access in ('write', 'pragma') else f[0, 0, 0] rhs = rcv[0] if access == 'read' else f[0, 0, 0] body = DummyExpr(lhs, rhs + 1) for i, size in reversed(list(enumerate(sizes))): pragmas = () if access == 'pragma' and i == len(sizes)-1: pragmas = (Pragma('unroll(1)'),) body = Iteration(body, Dimension(name=f'tap{i}'), (0, (size or n)-1, 1), pragmas=pragmas) body = ParlangTransformer._make_unroll_many_short(body) assert [str(p) for i in FindNodes(Iteration).visit(body) for p in i.pragmas] == ( [f'#pragma unroll({unroll})'] if unroll else [] ) @pytest.mark.parametrize('interpolation,r', [('linear', 1), ('sinc', 4)]) @pytest.mark.parametrize('dtype', [np.float32, np.float64]) @pytest.mark.parametrize('initial', [0., 2.5]) def test_indirect_reduction(self, interpolation, r, dtype, initial, host_kwargs): grid = Grid(shape=(97,), extent=(96.,), dtype=dtype) f = Function(name='f', grid=grid, space_order=r) rcv = SparseFunction(name='rcv', grid=grid, npoint=65, interpolation=interpolation, r=r) eqns = rcv.interpolate(expr=f, increment=True) op0 = Operator(eqns, name='IndirectReductionHost', **host_kwargs) op1 = Operator(eqns, name='IndirectReduction', opt=opt0({'gpu-reduce-local': False})) assert '#pragma unroll' not in str(op1) f.data_with_halo[:] = (np.arange(f.data_with_halo.size) % 7 - 3)/7. # Interior, partial and fully masked stencils at both ends positions = [0., r + .25, 3*r + .25, 48.25, 96.25 - 3*r, 96.25 - r, 96.] rcv.coordinates.data[:, 0] = np.resize(positions, rcv.npoint) bounds = dict(x_m=3*r, x_M=96 - 3*r) rcv.data[:] = initial op0.apply(**bounds) expected = rcv.data.copy() rcv.data[:] = initial op1.apply(**bounds) assert expected[0] == expected[6] == initial assert np.any(expected != initial) assert np.allclose(rcv.data, expected, rtol=1e-5, atol=1e-6) class TestWDR: def _check_gencode(self, op, expected, exp_shmas=None): kernel = op._func_table['kernel0'].root dims = FindSymbols('dimensions').visit(kernel) n_warpid = len([d for d in dims if isinstance(d, WarpId)]) if expected > 0: assert n_warpid == 1 warp_derivs = [i for i in op._func_table if i.startswith('warp_deriv')] assert len(warp_derivs) == expected else: assert n_warpid == 0 if exp_shmas is not None: functions = FindSymbols().visit(kernel) shmas = [f for f in functions if f._mem_shared] assert len(shmas) == exp_shmas def test_pickle_guard_rebuild(self): """ MFE for rebuilding a WDR Operator from a cloudpickle artifact. While reconstructing the surrounding Operator, SymPy may canonicalize this generated loader guard through ``as_dummy``. That must preserve the Devito ``Symbol`` used by ``FieldFromComposite``. """ grid = Grid(shape=(8, 8, 8)) z = grid.dimensions[-1] tz = ThreadDimLayoutA('tz', z, 0, z.symbolic_size - 1, 1, init_value=thread_idxs[0]) guard = Gt( z - tz + sympy.Mod(thread_idxs[0] + 32*thread_idxs[1], 32), z.symbolic_min - 8 ) guard = pickle.loads(pickle.dumps(guard)) # This is the SymPy reconstruction path reached while unpickling the # production WDR Operator. assert guard.as_dummy() == guard @pytest.mark.parametrize('par_tile,so,expected,exp_shmas', [ pytest.param((32, 4), 8, 4, 8, marks=skipif('device-hip', reason='tile too small')), pytest.param((32, 4), 16, 4, 7, marks=skipif('device-hip', reason='tile too small')), pytest.param((64, 4), 8, 6, 10, marks=skipif('device-hip', reason='runs out of shm')), pytest.param((64, 4), 16, 6, 8, marks=skipif(['device-hip', 'nocuda-cc80'], reason=('AMD GPU: mystery; ' 'Nvidia GPU cc<80: not enough shm'))), ]) @pytest.mark.parametrize('extra', [{}, {'gpu-prefetch': 4}]) def test_basic(self, par_tile, so, expected, exp_shmas, extra): grid = Grid(shape=(75, 38, 35)) nt = 5 f = Function(name='f', grid=grid, space_order=so) g = Function(name='g', grid=grid, space_order=so) u = TimeFunction(name='u', grid=grid, space_order=so, save=nt, layers=NoLayers) v0 = TimeFunction(name='v0', grid=grid, space_order=so) vs = [v0] + [v0.func(name=f'v{i}') for i in range(1, 6)] _, v1, v2, v3, v4, v5 = vs # noqa expr = v0.dy.dz + v1.dy.dz2 + v2.dy.dz3 + v3.dy.dz4 + v4.dy.dz5 + v5.dy.dz6 eqn = Eq(u.forward, expr + u + 1.) subs = grid.spacing_map op0 = Operator(eqn, subs=subs, opt=opt4({'par-tile': par_tile, 'gpu-fit': u})) op1 = Operator(eqn, subs=subs, opt=opt4({'errctl': 'basic', 'par-tile': par_tile, 'gpu-fit': u, 'gpu-wdr': 'auto', **extra})) # Check generated code # The reduced shared-array count locks in WDR reuse when a dead array # is available. Prefetched arrays remain reserved until their wait. self._check_gencode(op1, expected, exp_shmas=None if extra else exp_shmas) u1 = u.func(name='u1') f.data[:] = .01 g.data[:] = .02 for v in vs: init_chessboard(v) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-4) @skipif('device-hip') def test_codegen_odd(self): grid = Grid(shape=(75, 38, 35)) nt = 5 so = 8 u = TimeFunction(name='u', grid=grid, space_order=so, save=nt, layers=NoLayers) v0 = TimeFunction(name='v0', grid=grid, space_order=so) vs = [v0] + [v0.func(name=f'v{i}') for i in range(1, 6)] _, v1, v2, v3, v4, v5 = vs # noqa expr = v0.dy.dz + v1.dy.dz2 + v2.dy.dz3 + v3.dy.dz4 + v4.dy.dz5 + v5.dy.dz6 eqn = Eq(u.forward, expr + u + 1.) subs = grid.spacing_map op0 = Operator(eqn, name="op0", subs=subs, opt=opt4({'par-tile': (32, 4), 'gpu-fit': u})) op1 = Operator(eqn, name="op1", subs=subs, opt=opt4({'errctl': 'basic', 'par-tile': (32, 4), 'gpu-fit': u, 'gpu-wdr': 'auto'})) self._check_gencode(op1, 4) ccode = str(op1.ccode) assert 'template 0 assert '__launch_bounds__(512)' in ccode1 assert f'float q0_{so + 2} =' in ccode1 assert f'q0_0 + q0_{so - 1}' in ccode1 assert ccode1.count('switch (wid)') == 1 assert all(f'case {i}:' in ccode1 for i in range(8, 16)) assert ccode1.count('__syncthreads()') == ccode0.count('__syncthreads()') assert 'float wdr' not in ccode1 # Keep packed predicates out of the helper arguments, sharing the # mask across the three fields owned by each work unit calls = [i for i in FindNodes(Call).visit(op1._func_table['kernel0'].root) if i.name.startswith('warp_derivative_y')] masks = [] for call in calls: parameters = op1._func_table[call.name].root.parameters index = next(i for i, p in enumerate(parameters) if p.name == 'store_mask') masks.append(call.arguments[index]) assert len(masks) == 48 assert masks.count(15) == 24 masks = [i for i in masks if i != 15] assert all(i.is_Symbol and i.dtype == np.uint32 for i in masks) assert len(set(masks)) == 8 for source in sources: init_chessboard(source) overrides = {i.name: i.func(name=f'{i.name}1') for i in outputs} for field in (*outputs, *overrides.values()): field.data[:] = 100 op0.apply(time_M=nt-2, **bounds) op1.apply(time_M=nt-2, **bounds, **overrides) for field, other in zip(outputs, overrides.values()): assert np.allclose(field.data, other.data, rtol=1.e-3, atol=1.e-5) @skipif('device-hip') @pytest.mark.parametrize('access,before_barrier', [ (None, False), ('read', False), ('write', False), ('read', True), ('write', True) ]) def test_direct_store_movement(self, monkeypatch, access, before_barrier): from devitopro.passes import wdr from devitopro.passes.wdr import planner grid = Grid(shape=(17, 13, 35)) y = grid.dimensions[1] sources = [TimeFunction(name=f'u{i}', grid=grid, space_order=8, interp_order=8, staggered=y) for i in range(2)] outputs = [TimeFunction(name=f'v{i}', grid=grid, space_order=8, staggered=NODE) for i in range(2)] eqns = [Eq(v.forward, u) for u, v in zip(sources, outputs)] bundle_derivatives = wdr.bundle_derivatives enabled = False seen = [] def split_consumers(kernel, **kwargs): candidates = planner.collect_candidates(kernel, ('y',)) matches = [c for c in candidates if len(c.cluster.exprs) == 2 and all(e.lhs.function in outputs for e in c.cluster.exprs)] if not matches: return bundle_derivatives(kernel, **kwargs) candidate, = matches index = candidate.index cluster = candidate.cluster e0, e1 = cluster.exprs # Exercise motion across Clusters at the WDR input, after earlier # passes could fuse the assignments or remove the dead store clusters = list(kernel) clusters[index:index + 1] = [cluster.rebuild(exprs=(e,)) for e in (e0, e1)] if access is not None: expr = Eq(e0.lhs, e1.lhs) if access == 'read' else Eq(e1.lhs, 0) if before_barrier: # An access before the preceding barrier is outside the # movement interval and must not disable direct output insertion = max(i for i, c in enumerate(kernel[:index]) if c.is_thread_pool_sync) else: insertion = index + 1 clusters.insert(insertion, cluster.rebuild(exprs=(expr,))) kernel = kernel.rebuild(clusters=clusters) seen.append(enabled) return bundle_derivatives(kernel, **kwargs) if enabled else kernel monkeypatch.setattr(wdr, 'bundle_derivatives', split_consumers) options = optX({'par-tile': (64, 8, 48), 'gpu-fit': outputs, 'gpu-vect': 'stream-minimal', 'gpu-wdr': 'mz0:my1:p8', 'cire-hoist': 0, 'errctl': 'basic'}) op0 = Operator(eqns, name='op0', subs=grid.spacing_map, opt=options) enabled = True op1 = Operator(eqns, name='op1', subs=grid.spacing_map, opt=options) assert seen == [False, True] for source in sources: init_chessboard(source) for field in outputs: field.data[:] = 100 op0.apply(time_M=0) expected = [np.array(field.data) for field in outputs] for field in outputs: field.data[:] = 100 op1.apply(time_M=0) for field, reference in zip(outputs, expected): assert np.allclose(field.data, reference, rtol=1.e-5, atol=1.e-5) assert any(i.startswith('warp_derivative_y') for i in op1._func_table) ncaptures = 0 if access is None or before_barrier else 2 assert str(op1.ccode).count('float wdr') == ncaptures @skipif('device-hip') @pytest.mark.parametrize('dimension,phase_sizes,extra,gpu_wdr,expected', [ pytest.param('z', (1, 1), 2048, 'auto', 1, id='cumulative'), pytest.param('z', (1, 1), 4096, 'auto', 2, id='exact-fit'), pytest.param('z', (1, 1), 2047, 'auto', 0, id='below-one-tile'), pytest.param('z', (2,), 2048, 'auto', 0, id='single-phase-fallback'), pytest.param('z', (2, 1), 2048, 'auto', 1, id='rollback'), pytest.param('z', (2, 1), 6144, 'auto', 3, id='all-fit'), pytest.param('z', (1, 1), 2048, 'mz1:my0', None, id='explicit-cumulative'), pytest.param('z', (2,), 2048, 'mz1:my0', None, id='explicit-single'), pytest.param('y', (1, 1), 0, 'auto', 2, id='direct-no-scratch'), ]) def test_scratch_budget(self, monkeypatch, dimension, phase_sizes, extra, gpu_wdr, expected): from copy import copy from devitopro.passes import wdr from devitopro.passes.shm import shmas_nbytes from devitopro.passes.wdr import planner grid = Grid(shape=(17, 13, 35), extent=(16., 12., 34.)) d = grid.dimensions[-1 if dimension == 'z' else -2] u = TimeFunction(name='u', grid=grid, space_order=8) outputs = [TimeFunction(name=f'v{i}', grid=grid, space_order=8) for i in range(sum(phase_sizes))] eqns = [Eq(v.forward, Derivative(u, (d, i))) for i, v in enumerate(outputs, 1)] bundle_derivatives = wdr.bundle_derivatives enabled = False seen = [] def split_phases(kernel, **kwargs): candidates = planner.collect_candidates(kernel, (dimension,)) matches = [c for c in candidates if len(c.cluster.exprs) == len(outputs) and all(e.lhs.function in outputs for e in c.cluster.exprs)] if not matches: return bundle_derivatives(kernel, **kwargs) candidate, = matches index = candidate.index cluster = candidate.cluster # Reload the shared input between consumers to form distinct WDR # phases. All derivatives read that array, so it cannot be reused # for the hand-off; each result needs a dedicated 8x64 float tile clusters = list(kernel[:index]) start = 0 for size in phase_sizes: if start: clusters.extend(kernel[:index]) clusters.append(cluster.rebuild(exprs=cluster.exprs[start:start + size])) start += size clusters.extend(kernel[index + 1:]) kernel = kernel.rebuild(clusters=clusters) candidates = planner.collect_candidates(kernel, (dimension,)) assert len(planner.build_phases(candidates, kernel)) == len(phase_sizes) seen.append(enabled) if not enabled: return kernel # Limit only WDR's budget, leaving input allocation and the actual # device limits unchanged. This makes the boundary deterministic platform = copy(kwargs['platform']) base = shmas_nbytes(kernel.shmas) platform.max_shm_per_block = base + extra kwargs['platform'] = platform result = bundle_derivatives(kernel, **kwargs) assert shmas_nbytes(result.shmas) <= platform.max_shm_per_block return result monkeypatch.setattr(wdr, 'bundle_derivatives', split_phases) options = optX({'par-tile': (64, 8, 48), 'gpu-fit': outputs, 'gpu-vect': 'stream-minimal', 'gpu-wdr': gpu_wdr, 'gpu-prefetch': False, 'gpu-shm': 'min-reuse', 'cire-hoist': 0, 'errctl': 'basic'}) op0 = Operator(eqns, name='op0', subs=grid.spacing_map, opt=options) enabled = True if expected is None: with pytest.raises(InvalidOperator, match='WDR exceeded the shared-memory'): Operator(eqns, name='op1', subs=grid.spacing_map, opt=options) assert seen == [False, True] return op1 = Operator(eqns, name='op1', subs=grid.spacing_map, opt=options) assert seen == [False, True] self._check_gencode(op1, expected) init_chessboard(u) overrides = {v.name: v.func(name=f'{v.name}1') for v in outputs} for field in (*outputs, *overrides.values()): field.data[:] = 100 op0.apply(time_M=0) op1.apply(time_M=0, **overrides) for field, other in zip(outputs, overrides.values()): assert np.allclose(field.data, other.data, rtol=1.e-4, atol=1.e-5) @skipif('device-hip') def test_direct_fallback_on_later_read(self): grid = Grid(shape=(50, 33, 67)) nt = 5 so = 8 y = grid.dimensions[-2] sources = [TimeFunction(name=f'u{i}', grid=grid, space_order=so, interp_order=so, staggered=y) for i in range(6)] outputs = [TimeFunction(name=f'v{i}', grid=grid, space_order=so, staggered=NODE) for i in range(6)] w = TimeFunction(name='w', grid=grid, space_order=so) eqns = [Eq(v.forward, u) for u, v in zip(sources, outputs)] eqns.append(Eq(w.forward, sum(v.forward for v in outputs))) fields = (*outputs, w) options = {'par-tile': (64, 8, 48), 'gpu-fit': fields} op0 = Operator(eqns, name='op0', subs=grid.spacing_map, opt=opt4(options)) op1 = Operator(eqns, name='op1', subs=grid.spacing_map, opt=optX({**options, 'errctl': 'basic', 'cire-hoist': 0, 'gpu-vect': 'stream-minimal', 'gpu-wdr': 'mz0:my3:p8'})) assert any(i.startswith('warp_derivative_y') for i in op1._func_table) assert 'float wdr' in str(op1.ccode) for source in sources: init_chessboard(source) overrides = {i.name: i.func(name=f'{i.name}1') for i in fields} for field in (*fields, *overrides.values()): field.data[:] = 100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) for field, other in zip(fields, overrides.values()): assert np.allclose(field.data, other.data, rtol=1.e-3, atol=1.e-5) @skipif('device-hip') def test_auto_bundle_prefix(self): grid = Grid(shape=(50, 33, 67)) nt = 5 y = grid.dimensions[-2] sources = [TimeFunction(name=f'u{i}', grid=grid, space_order=8, interp_order=8, staggered=y) for i in range(9)] outputs = [TimeFunction(name=f'v{i}', grid=grid, space_order=8, staggered=NODE) for i in range(9)] eqns = [Eq(v.forward, u) for u, v in zip(sources, outputs)] options = {'par-tile': (64, 4), 'gpu-fit': outputs} op0 = Operator(eqns, name='op0', subs=grid.spacing_map, opt=opt4(options)) op1 = Operator(eqns, name='op1', subs=grid.spacing_map, opt=optX({**options, 'errctl': 'basic', 'cire-hoist': 0, 'gpu-vect': 'stream-minimal', 'gpu-wdr': 'auto'})) # Three M=3 bundles over two subtiles create six work units. Four # producer warps can balance only a two-bundle prefix; the remaining # bundle must stay on the ordinary consumer path. ccode = str(op1.ccode) assert ccode.count('switch (wid)') == 1 assert ccode.count('case ') == 4 assert ccode.count('float wdr') == 6 for source in sources: init_chessboard(source) overrides = {i.name: i.func(name=f'{i.name}1') for i in outputs} for field in (*outputs, *overrides.values()): field.data[:] = 100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) for field, other in zip(outputs, overrides.values()): assert np.allclose(field.data, other.data, rtol=1.e-3, atol=1.e-5) @skipif('nocuda') @switchconfig(platform='ampere') def test_factored_nested_derivs_explicit(self): grid = Grid(shape=(32, 32, 32), extent=(630., 630., 630.)) so = (16, 8, 8) w = TimeFunction(name='w', grid=grid, time_order=2, space_order=so, save=Buffer(2)) u = Function(name='u', grid=grid, space_order=so) v = Function(name='v', grid=grid, space_order=so) expr = w.dx2 + w.dy2 + w.dz2 + v.dx + v.dz + u.dz eq = w.dt2 - expr eqns = [Eq(u, w.dx), Eq(v, w.dy), Eq(w.forward, solve(eq + w.dt, w.forward))] op0 = Operator(eqns, name='op0', opt=optX({'errctl': 'basic', 'gpu-fit': [w], 'gpu-prefetch': 4, 'gpu-vect': False, 'par-tile': (64, 4, 128)})) op1 = Operator(eqns, name='op1', opt=optX({'errctl': 'basic', 'gpu-fit': [w], 'gpu-prefetch': 4, 'gpu-vect': False, 'gpu-wdr': 'auto', 'par-tile': (64, 4, 128)})) w1 = w.func(name='w1') for i in (w, w1): i.data.fill(0.) mid = tuple(n // 2 for n in grid.shape) i.data[(slice(None), *mid)] = 1. op0.apply(time_M=1, dt=1.) op1.apply(time_M=1, dt=1., w=w1) assert np.allclose(w.data, w1.data, rtol=1.e-5, atol=1.e-6) @skipif('nocuda') @switchconfig(platform='ampere') def test_nested_dz_w_prefetch(self): grid = Grid(shape=(75, 38, 35)) nt = 5 so = 8 u = TimeFunction(name='u', grid=grid, space_order=so, save=nt, layers=NoLayers) vs = [TimeFunction(name=f'v{i}', grid=grid, space_order=so) for i in range(4)] eqn = Eq(u.forward, sum(v.dy.dz for v in vs) + u + 1.) subs = grid.spacing_map op0 = Operator(eqn, name="op0", subs=subs, opt=opt4({'par-tile': (32, 4), 'gpu-fit': u})) op1 = Operator(eqn, name="op1", subs=subs, opt=opt4({'errctl': 'basic', 'par-tile': (32, 4), 'gpu-fit': u, 'gpu-prefetch': 1, 'gpu-wdr': 'auto'})) self._check_gencode(op1, 1) u1 = u.func(name='u1') for v in vs: init_chessboard(v) init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-3, atol=1.e-3) @pytest.mark.parametrize('par_tile', [ pytest.param((16, 8), marks=skipif('device-hip', reason='tile too small')), pytest.param((32, 8), marks=skipif('device-hip', reason='tile too small')), (64, 8), ]) def test_etti_velocity_like(self, par_tile): grid = Grid(shape=(72, 15, 31)) x, y, z = grid.dimensions nt = 5 f = Function(name='f', grid=grid, space_order=16) v = VectorTimeFunction(name="v", grid=grid, space_order=16, save=nt, layers=NoLayers) tau = TensorTimeFunction(name="tau", grid=grid, space_order=16) fields = tuple(v) expr = v.dt - f * div(tau) eqns = [Eq(v.forward, solve(expr, v.forward))] op0 = Operator(eqns, opt=opt4({'gpu-fit': fields, 'par-tile': par_tile})) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'cire-hoist': 0, # Legacy 'gpu-fit': fields, 'par-tile': par_tile, 'gpu-wdr': 'auto'})) # Check generated code. The number of spatial work units depends on the # platform's thread-group size, so assert the Dimension coverage. warp_derivs = [i for i in op1._func_table if i.startswith('warp_derivative')] assert any(i.startswith('warp_derivative_z') for i in warp_derivs) assert any(i.startswith('warp_derivative_y') for i in warp_derivs) f.data[:] = 3.5 tau[0].data[:] = 5.0 tau[1].data[:] = 6.2 tau[2].data[:] = 7.3 tau[4].data[:] = 1.2 tau[5].data[:] = 9.3 tau[8].data[:] = 3.4 overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} for objs in (fields, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2, dt=.001) op1.apply(time_M=nt-2, dt=.001, **overrides) for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) @skipif('nocuda') @switchconfig(platform='ampere') def test_w_prefetch(self): grid = Grid(shape=(50, 33, 41)) so = 8 nt = 5 f = Function(name='f', grid=grid, space_order=so) v = VectorTimeFunction(name="v", grid=grid, space_order=so) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so, save=nt, layers=NoLayers) vx, _, vz = v tau_xz = tau[2] eq = Eq(tau_xz.forward, vx.dx + vz.dz + f*vx.dz) op0 = Operator(eq, opt=opt4({'gpu-fit': tau_xz, 'par-tile': (32, 4, 48), 'gpu-flip-derivs': True})) op1 = Operator(eq, opt=opt4({'errctl': 'basic', 'cire-hoist': 0, # Legacy 'gpu-fit': tau_xz, 'par-tile': (32, 4, 48), 'gpu-flip-derivs': True, 'gpu-prefetch': True, 'gpu-wdr': 'auto'})) # Check generated code self._check_gencode(op1, 3) init_chessboard(f) init_chessboard(vx) init_chessboard(vz) tau_xz1 = tau_xz.func(name='tau_xz1') tau_xz.data[:] = 100 tau_xz1.data[:] = 100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, tau_xz=tau_xz1) assert np.allclose(tau_xz.data, tau_xz1.data, rtol=1.e-3, atol=1.e-5) @skipif('nocuda') @switchconfig(platform='ampere') @pytest.mark.parametrize('par_tile,expected,exp_shmas,extra', [ ((24, 4, 32), 0, 35, {}), ((32, 4, 48), 3, 60, {}), pytest.param((64, 4, 48), 4, 64, {}, marks=skipif('nocuda-cc80')), # The following is to ensure the par-tile is not too big for Volta (cc=70), # but we have to disable prefetching and use min-reuse to avoid running # out of memory or in proximity of the shared memory capacity ((64, 4, 48), 4, 64, {'gpu-prefetch': False, 'gpu-shm': 'min-reuse'}), ]) def test_etti_stress_like(self, par_tile, expected, exp_shmas, extra): grid = Grid(shape=(50, 33, 41)) so = 8 nt = 5 f = Function(name='f', grid=grid, space_order=so) g = Function(name='g', grid=grid, space_order=so) v = VectorTimeFunction(name="v", grid=grid, space_order=so) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so, save=nt, layers=NoLayers) vx, vy, vz = v expr = (vx.dx + vy.dy + vz.dz + f*(vx.dy + vy.dx) + g*(vx.dz + vz.dx) + f*(vy.dz + vz.dy)) fields = tau.values() eqns = [Eq(i.forward, expr) for i in fields] op0 = Operator(eqns, opt=opt4({'gpu-fit': fields, 'par-tile': par_tile, 'gpu-flip-derivs': True})) op1 = Operator(eqns, opt=optX({'errctl': 'basic', 'cire-hoist': 0, # Legacy 'gpu-fit': fields, 'par-tile': par_tile, 'gpu-flip-derivs': True, 'gpu-vect': 'stream-minimal', 'gpu-wdr': 'auto', **extra})) self._check_gencode(op1, expected, exp_shmas=exp_shmas) f.data[:] = 19.0 g.data[:] = 18.0 vx.data[:] = 11.0 vy.data[:] = 12.0 vz.data[:] = 7.0 overrides = {i.name: i.func(name='%s1' % i.name) for i in fields} for objs in (fields, overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, **overrides) for i, j in zip(fields, overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3, atol=1.e-3) def test_eiso_like(self): grid = Grid(shape=(60, 60, 60)) so = (16, 8, 8) v = VectorTimeFunction(name="v", grid=grid, space_order=so) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so) mu = Function(name='mu', grid=grid, space_order=so, avg_mode='safe_harmonic') e = (grad(v) + grad(v).transpose(inner=False)) expr_tau = tau.dt - diag(div(v)) - mu*e eq = Eq(tau.forward, solve(expr_tau, tau.forward)) op0 = Operator(eq, opt=opt4({'par-tile': (32, 8), 'gpu-flip-derivs': True})) op1 = Operator(eq, opt=optX({'errctl': 'basic', 'gpu-flip-derivs': True, 'gpu-vect': 'stream-minimal', 'gpu-prefetch': (False, 4), 'par-tile': ((16, 8), (32, 8, 16)), 'gpu-wdr': (False, 'auto')})) mu.data[:] = 3.5 vx, vy, vz = v vx.data[:] = 11.0 vy.data[:] = 12.0 vz.data[:] = 7.0 overrides = {i.name: i.func(name=f'{i.name}1') for i in tau.values()} for objs in (tau.values(), overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(dt=0.001, time_M=1) op1.apply(dt=0.001, time_M=1, **overrides) for i, j in zip(tau.values(), overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) @skipif('nocuda') @pytest.mark.parametrize('dimension', ['y', 'z']) @pytest.mark.parametrize('symmetric', [False, True]) @pytest.mark.parametrize('nfactors,shifted,padding', [ (1, False, padding) for padding in ((1, 0), (0, 1), (2, 1)) ] + list(product((2, 3), (False, True), ((0, 0), (1, 0), (0, 1))))) def test_padded_and_product_inputs(self, dimension, symmetric, nfactors, shifted, padding): grid = Grid(shape=(25, 19, 37), extent=(24., 18., 36.)) fields = [TimeFunction(name=f'u{i}', grid=grid, space_order=8) for i in range(nfactors)] v = TimeFunction(name='v', grid=grid, space_order=8) expr = fields[0] y, z = grid.dimensions[1:] for i, f in enumerate(fields[1:], 1): expr *= f.shift(y, i*y.spacing).shift(z, -i*z.spacing) if shifted else f left, right = padding coeffs = [1., -3.] if not any(padding): coeffs += [2. if symmetric else 0.] coeffs += [-3., 1.] if symmetric else [3., -1.] weights = [0.] * left + coeffs + [0.] * right deriv = Derivative(expr, {'y': y, 'z': z}[dimension], weights=weights, fd_order=len(weights)-1) eqn = Eq(v.forward, v + deriv) options = {'par-tile': (64, 8, 32), 'gpu-vect': 'stream-only', 'cire-mode': '1pass', 'cire-hoist': 0, 'gpu-prefetch': 4} reference = Operator(eqn, name='reference', subs=grid.spacing_map, opt=opt4(options)) mode = 'mz0:my1:p4' if dimension == 'y' else 'mz1:my0:p4' candidate = Operator(eqn, name='candidate', subs=grid.spacing_map, opt=opt4({**options, 'gpu-wdr': mode})) self._check_gencode(candidate, 1) # Dyadic inputs and unit spacing make the indexing check exact in float32. rng = np.random.default_rng(42) scale = 16 if nfactors == 1 else 4 for f in [*fields, v]: values = rng.integers(-scale, scale+1, f.data_with_halo.shape) f.data_with_halo[:] = values/scale v1 = v.func(name='v1') v1.data_with_halo[:] = v.data_with_halo reference.apply(time_M=3) candidate.apply(time_M=3, v=v1) np.testing.assert_array_equal(v1.data, v.data) @skipif('nocuda') @pytest.mark.parametrize('dimension', ['y', 'z']) @pytest.mark.parametrize('weights,dtype,stride,expected', [ ([0., 1., 0., -1., 0.], np.float32, 1, 1), ([0., 0., 0.], np.float32, 1, 1), ([0., 1., 2., 3., 0.], np.float32, 1, 0), ([1., -3., 0., 3., -1.], np.float64, 1, 0), ([1., -3., 0., 3., -1.], np.float32, 2, 1), ], ids=['symmetric_padding', 'zero', 'unpaired', 'mixed', 'strided']) def test_product_edge_cases(self, dimension, weights, dtype, stride, expected): grid = Grid(shape=(25, 19, 37), extent=(24., 18., 36.)) a = TimeFunction(name='a', grid=grid, space_order=8) b = TimeFunction(name='b', grid=grid, space_order=8, dtype=dtype) v = TimeFunction(name='v', grid=grid, space_order=8) axis = 1 if dimension == 'y' else 2 d = grid.dimensions[axis] expr = a*b.subs(d, stride*d) eqn = Eq(v.forward, v + Derivative(expr, d, weights=weights, fd_order=len(weights)-1)) options = {'par-tile': (64, 8, 32), 'gpu-vect': 'stream-only', 'cire-mode': '1pass', 'cire-hoist': 0, 'gpu-prefetch': 4} reference = Operator(eqn, name='reference', subs=grid.spacing_map, opt=opt4(options)) candidate = Operator(eqn, name='candidate', subs=grid.spacing_map, opt=opt4({**options, 'gpu-wdr': 'auto'})) self._check_gencode(candidate, expected) rng = np.random.default_rng(42) for f in (a, b, v): f.data_with_halo[:] = rng.integers(-4, 5, f.data_with_halo.shape)/4 v1 = v.func(name='v1') v1.data_with_halo[:] = v.data_with_halo # Keep doubled source indices inside the allocated domain bounds = {f'{dimension}_M': grid.shape[axis]//stride-1} reference.apply(time_M=3, **bounds) candidate.apply(time_M=3, v=v1, **bounds) np.testing.assert_array_equal(v1.data, v.data) @skipif('nocuda') @pytest.mark.parametrize('dimension', ['y', 'z']) @pytest.mark.parametrize('reuse', ['min-reuse', 'max-reuse']) def test_mixed_product_bundle(self, dimension, reuse): grid = Grid(shape=(25, 19, 37), extent=(24., 18., 36.)) fields = [TimeFunction(name=f'u{i}', grid=grid, space_order=8) for i in range(3)] a, b, c = fields v = TimeFunction(name='v', grid=grid, space_order=8) d = grid.dimensions[1 if dimension == 'y' else 2] derivs = [Derivative(expr, d, weights=[1., -3., 0., 3., -1.], fd_order=4) for expr in (a, a*b, a*b*c)] # Mix atomic and product derivatives while every factor evolves eqns = [Eq(v.forward, v + sum(derivs))] eqns += [Eq(f.forward, f + (i+1)/4) for i, f in enumerate(fields)] options = {'par-tile': (64, 8, 32), 'gpu-vect': 'stream-only', 'cire-mode': '1pass', 'cire-hoist': 0, 'gpu-prefetch': 4, 'gpu-shm': reuse} reference = Operator(eqns, name='reference', subs=grid.spacing_map, opt=opt4(options)) mode = 'mz0:my3:p4' if dimension == 'y' else 'mz3:my0:p4' candidate = Operator(eqns, name='candidate', subs=grid.spacing_map, opt=opt4({**options, 'gpu-wdr': mode})) self._check_gencode(candidate, 1) rng = np.random.default_rng(42) outputs = [*fields, v] overrides = {f.name: f.func(name=f'{f.name}1') for f in outputs} for f in outputs: f.data_with_halo[:] = rng.integers(-4, 5, f.data_with_halo.shape)/4 overrides[f.name].data_with_halo[:] = f.data_with_halo reference.apply(time_M=3) candidate.apply(time_M=3, **overrides) for f in outputs: np.testing.assert_array_equal(overrides[f.name].data, f.data) class TestKernelOptRegRotAdvanced: def _retrieve_x_indices(self, op, function): kernel = op._func_table['kernel0'].root indexeds = FindSymbols('indexeds').visit(kernel) indexeds = [i for i in indexeds if i.function is function] indices = [str(i.indices[:2]) for i in indexeds] return indices def _check_scattering(self, op, flag=True, kid='kernel0', nqueues=None, expected=None): # Pretty rudimentary, but it does the job kernel = op._func_table[kid].root indexeds = FindSymbols('indexeds').visit(kernel) queues = [i for i in indexeds if isinstance(i.function, TBQueue)] test = any(f"{i.name}[j] +=" in str(kernel) for i in queues) assert flag == test if nqueues is not None: assert len(set(i.function for i in queues)) == nqueues if expected is not None: assert isinstance(expected, dict) for k, v in expected.items(): found = [i for i in indexeds if i.name == k] assert len(found) == v @pytest.mark.parametrize('opt,expected', [ (opt4({'par-tile': (4, 4, 4), 'gpu-regrot': 'scatter'}), (11, 0, 5, 0, 5, 0, 1)), (opt4({'par-tile': (4, 4, 4), 'gpu-regrot': 'scatter', 'gpu-glb-mem': False, 'gpu-prefetch': True}), (21, 5, 0, 5, 0, 5, 1)), ]) def test_laplace(self, opt, expected): grid = Grid(shape=(16, 16, 16)) u = TimeFunction(name='u', grid=grid, space_order=4) u1 = u.func(name='u1') eq = Eq(u.forward, u.laplace + .01) op0 = Operator(eq, opt=opt4({'par-tile': (4, 4, 4)})) op1 = Operator(eq, opt=opt) # Check generated code # All stencil accesses along x (or qq0, for initialization) are expected # to be at the same exact offset (e.g., x+6) indices = self._retrieve_x_indices(op1, u) e0, e1, e2, e3, e4, e5, e6 = expected assert len(indices) == e0 assert len([i for i in indices if i == '(t0, x0_blk0 + qq0 + 3)']) == e1 assert len([i for i in indices if i == '(t0, x0_blk0 + qq0 + 2)']) == e2 assert len([i for i in indices if i == '(t0, x + 7)']) == e3 assert len([i for i in indices if i == '(t0, x + 6)']) == e4 assert len([i for i in indices if i == '(t0, x0_blk0 + 6)']) == e5 assert len([i for i in indices if i == '(t1, x + 4)']) == e6 # Write out op0.apply(time_M=3) op1.apply(time_M=3, u=u1) expected = np.linalg.norm(u.data) assert expected > 0 # NOTE: Not exact equality because of the different order of operations assert np.isclose(expected, np.linalg.norm(u1.data), rtol=1.e-6) @pytest.mark.parametrize('pde,nqueues,nindices', [ ('m * u.dt2 - u.laplace', 2, 5), ('(1 / (vp*vp)) * u.dt2 - u.laplace', 2, 5), ('b + (1 / (vp*vp)) * u.dt2 - u.laplace', 2, 5), ('(b / (vp*vp)) * u.dt2 - b*u.laplace', 2, 4), ]) def test_acoustic_iso_like(self, pde, nqueues, nindices, extra=None): extra = extra or {} grid = Grid(shape=(50, 50, 50)) b = Constant(name='b') m = Function(name='m', grid=grid, space_order=8) vp = Function(name='vp', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, space_order=8, time_order=2) pde = eval(pde) eq = Eq(u.forward, solve(pde, u.forward)) op0 = Operator(eq, opt=opt4({'par-tile': (8, 8, 16)})) op1 = Operator(eq, opt=opt4({'errctl': 'basic', 'par-tile': (8, 8, 16), 'gpu-regrot': 'scatter', **extra})) # Check generated code self._check_scattering(op1, nqueues=nqueues) indices = self._retrieve_x_indices(op1, u) assert len(set(indices)) == nindices if extra.get('gpu-glb-mem'): assert 't0' not in str(indices) else: assert '(t0, x + 12)' in indices assert '(t0, x + 8)' not in indices u1 = u.func(name='u1') b.data = 1.0 m.data[:] = .01 vp.data[:] = 10.0 init_chessboard(u) init_chessboard(u1) op0.apply(time_M=3, dt=.01) op1.apply(time_M=3, dt=.01, u=u1) assert np.linalg.norm(u.data) > 0 assert np.allclose(u.data, u1.data, rtol=1.e-4) @switchconfig(ignore_unknowns=True) @pytest.mark.parametrize('component', [ 0, None, ]) def test_divergence(self, component, space_order=8, optoptions=None): optoptions = optoptions or {} grid = Grid(shape=(20, 20, 20)) v = VectorTimeFunction(name='v', grid=grid, space_order=space_order) tau = TensorTimeFunction(name='tau', grid=grid, space_order=space_order) if component is None: eqn = Eq(v.forward, v - div(tau)) else: assert component in range(0, 3) eqn = Eq(v.forward[component], v[component] - div(tau)[component]) op0 = Operator(eqn, opt=opt4({'par-tile': (8, 8, 4)})) op1 = Operator(eqn, opt=opt4({'errctl': 'basic', 'par-tile': (8, 8, 4), 'gpu-regrot': 'scatter', **optoptions})) # Check generated code if component is None: if optoptions.get('gpu-glb-mem') == 'tma': expected = {} else: expected = {'tau_xy': 6, 'tau_xz': 6} self._check_scattering(op1, nqueues=3, expected=expected) else: # With just one component, this use case falls back to plain # register queues, since `summand_in = 0` given that all fronts # are already perfectly aligned self._check_scattering(op1, nqueues=1, flag=False) v_x, v_y, v_z = {i.name: i.func(name=f'{i.name}1') for i in v.values()}.values() for i in tau.values(): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, v_x=v_x, v_y=v_y, v_z=v_z) if component is not None: assert np.linalg.norm(v[component].data) > 0 for i, j in zip(v.values(), [v_x, v_y, v_z]): assert np.allclose(i.data, j.data, rtol=1.e-5) def test_avoid_scattering_if_isolated(self, gpu_glb_mem=False): grid = Grid(shape=(10, 10, 10)) u0 = TimeFunction(name='u0', grid=grid, space_order=4) u1 = u0.func(name='u1') v = u0.func(name='v') w = u0.func(name='w') eqns = [Eq(u0.forward, u0 - v.dx + v.dy + w.dy), Eq(u1.forward, u1 - w)] op0 = Operator(eqns, opt=opt4({'par-tile': (8, 8, 4)})) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': (8, 8, 4), 'gpu-regrot': 'scatter', 'gpu-glb-mem': gpu_glb_mem})) # Check generated code -- ensure we're using the scatter mode only # for the `v.dx + v.dy` fragment indices = self._retrieve_x_indices(op1, v) assert len(indices) == 6 assert len([i for i in indices if i == '(t0, x0_blk0 + qq0 + 2)']) == 3 assert len([i for i in indices if i == '(t0, x + 6)']) == 3 indices = self._retrieve_x_indices(op1, w) assert len(indices) == 3 assert all(i == '(t0, x + 4)' for i in indices) indices = self._retrieve_x_indices(op1, u0) assert len(indices) == 2 assert len([i for i in indices if i == '(t0, x + 4)']) == 1 assert len([i for i in indices if i == '(t1, x + 4)']) == 1 u0_B = u0.func(name='u0_B') u1_B = u1.func(name='u1_B') for i in (u0, u1, u0_B, u1_B, v, w): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, u0=u0_B, u1=u1_B) assert np.allclose(u0.data, u0_B.data, rtol=1.e-5) assert np.allclose(u1.data, u1_B.data, rtol=1.e-5) def test_avoid_scattering_if_not_Add(self, gpu_glb_mem=False): grid = Grid(shape=(20, 20, 20)) f = Function(name='f', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, space_order=8) eqn = Eq(u.forward, u.dx + f*u.dy.dx + u.dx.dz + u.dy.dz) op0 = Operator(eqn, opt=opt4({'par-tile': (8, 8, 4)})) op1 = Operator(eqn, opt=opt4({'errctl': 'basic', 'par-tile': (8, 8, 4), 'gpu-regrot': 'scatter', 'gpu-glb-mem': gpu_glb_mem})) # Check generated code self._check_scattering(op1, flag=False) u1 = u.func(name='u1') f.data[:] = 3.5 for i in (u, u1): init_chessboard(i) op0.apply(time_M=3) op1.apply(time_M=3, u=u1) assert np.linalg.norm(u.data) > 0 assert np.allclose(u.data, u1.data, rtol=1.e-4) @skipif('device-hip') def test_etti_velocity_like(self, gpu_glb_mem=False): TestAsyncPipeline().run_etti_velocity_like( 1, {'gpu-regrot': 'scatter', 'gpu-vect': False, 'gpu-glb-mem': gpu_glb_mem}, mode=1, expected=(0, 0, 0, 2) ) @skipif('device-hip') @switchconfig(ignore_unknowns=True) @pytest.mark.parametrize('exprs,expected,tol', [ (('vz.dx', 'vz.dy'), {'v_z': 3}, None), (('g*vz.dx+f*vz.dy', 'g*vz.dx+f*vz.dy'), {'v_z': 4, 'f': 5, 'g': 5}, None), (('vz.dz+g*vz.dx+f*vz.dy', 'vz.dz+g*vz.dx+f*vz.dy'), {'v_z': 4, 'f': 5, 'g': 5}, None), (('vz.dx', 'vz.dx', 'vz.dx'), {'v_z': 3}, None), (('1.', 'vx.dz + vy.dx'), {'v_x': 1, 'v_y': 2}, None), (('1.', 'vx.dz + g*vy.dx'), {'v_x': 1, 'v_y': 2, 'g': 4}, None), (('vx + f*vy.dx + 1', 'vx', 'f*vx.dz'), {'v_x': 2, 'v_y': 2, 'f': 2}, None), (('vx.dz + vx.dy + f*vy.dx + g*vz.dx + C + 1', 'f*vz.dx + vx.dz + g*vy.dx + vx.dy + C + 1', 'f*vy.dx + vx.dy + g*vz.dx + vx.dz + C + 1', 'f*vx.dy + vy.dx + g*vz.dy + vy.dz + C + 1', 'f*vx.dz + g*vy.dz + vz.dy + vz.dx + C + 1'), {'v_x': 2, 'v_y': 2, 'v_z': 3, 'f': 4, 'g': 4}, None), (('vx.dx + vx.dz + vx.dy + f*vy.dx + g*vz.dx + fields[0] + 1', 'vx.dx + f*vz.dx + vx.dz + g*vy.dx + vx.dy + fields[1] + 1', 'vx.dx + f*vy.dx + vx.dy + g*vz.dx + vx.dz + fields[2] + 1', 'vy.dy + f*vx.dy + vy.dz + vy.dx + g*vz.dy + fields[3] + 1', 'vy.dy + f*vx.dy + vy.dx + g*vz.dy + vy.dz + fields[4] + 1', 'vz.dz + f*vx.dz + g*vy.dz + vz.dy + vz.dx + fields[5] + 1'), {'v_x': 2, 'v_y': 3, 'v_z': 3, 'f': 4, 'g': 4}, 1e-1), ]) def test_etti_stress_like_incremental(self, exprs, expected, tol, gpu_glb_mem=False): _, op1 = TestAsyncPipeline().run_etti_stress_like_incremental( exprs, 0, 0, 0, None, {'gpu-prefetch': False, 'gpu-regrot': 'scatter', 'gpu-glb-mem': gpu_glb_mem}, tol=tol ) # Check generated code # We do NOT expect scattering in any of these tests (hence `flag=False`) # however we expect one single front per Function in the x-loop. Hence # the check on the indices self._check_scattering(op1, flag=False) for k, v in expected.items(): f, = [i for i in op1.parameters if i.name == k] indices = self._retrieve_x_indices(op1, f) assert len(set(indices)) == v def run_eiso_vel_like(self, extra, nqueues): grid = Grid(shape=(42, 15, 31)) x, y, z = grid.dimensions f = Function(name='f', grid=grid, space_order=8) m = Function(name='m', grid=grid, space_order=8, avg_mode='safe_harmonic') v = VectorTimeFunction(name="v", grid=grid, space_order=8) tau = TensorTimeFunction(name="tau", grid=grid, space_order=8) expr = v.dt - f * div(tau) eq = Eq(v.forward, solve(expr, v.forward)) mmask = INT(Gt(m, Cast(0, dtype=m.dtype))) sheartau = {t for t in tau.values() if t not in tau.diagonal()} eq = eq.subs({der: mmask * der for der in retrieve_derivatives(eq) if der.expr in sheartau}) op0 = Operator(eq, opt=opt4({'par-tile': (32, 4), 'gpu-flip-derivs': True})) op1 = Operator(eq, opt=opt4({'errctl': 'basic', 'cire-hoist': 0, # Legacy 'gpu-flip-derivs': True, 'par-tile': ((16, 4), (32, 4, 16)), 'gpu-regrot': 'scatter', **extra})) self._check_scattering(op1, nqueues=nqueues) f.data[:] = .3 m.data[:] = 3.5 txx, txy, txz, tyy, tyz, tzz = tau.values() txx.data[:] = 11.0 txy.data[:] = 12.0 txz.data[:] = 7.0 tyy.data[:] = 9.2 tyz.data[:] = 3.1 tzz.data[:] = 17.5 overrides = {i.name: i.func(name=f'{i.name}1') for i in v.values()} for objs in (v.values(), overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(dt=0.001, time_M=3) op1.apply(dt=0.001, time_M=3, **overrides) for i, j in zip(v.values(), overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) return op1 @skipif('device-hip') @pytest.mark.parametrize('extra,nqueues', [ ({}, 5), ({'gpu-prefetch': 4, 'gpu-glb-mem': 'async-pipe'}, 7), ]) def test_eiso_vel_like(self, extra, nqueues): self.run_eiso_vel_like(extra, nqueues) def run_eiso_stress_like(self, extra, nqueues): grid = Grid(shape=(60, 60, 60)) so = (16, 8, 8) v = VectorTimeFunction(name="v", grid=grid, space_order=so) tau = TensorTimeFunction(name="tau", grid=grid, space_order=so) damp = Function(name='damp', grid=grid, space_order=so) lam = Function(name='lam', grid=grid, space_order=so) mu = Function(name='mu', grid=grid, space_order=so, avg_mode='safe_harmonic') e = (grad(v) + grad(v).transpose(inner=False)) expr_tau = tau.dt - lam*diag(div(v)) - mu*e + damp*(tau + tau.forward) eqns = Eq(tau.forward, solve(expr_tau, tau.forward)) op0 = Operator(eqns, opt=opt4({'par-tile': (32, 8), 'gpu-flip-derivs': True})) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'gpu-flip-derivs': True, 'gpu-regrot': 'scatter', 'par-tile': ((16, 8), (32, 8, 16)), **extra})) self._check_scattering(op1, nqueues=nqueues, kid='kernel1') mu.data[:] = 3.5 vx, vy, vz = v vx.data[:] = 11.0 vy.data[:] = 12.0 vz.data[:] = 7.0 overrides = {i.name: i.func(name=f'{i.name}1') for i in tau.values()} for objs in (tau.values(), overrides.values()): for i, ff in enumerate(objs, 1): ff.data[:] = i*100 op0.apply(dt=0.001, time_M=3) op1.apply(dt=0.001, time_M=3, **overrides) for i, j in zip(tau.values(), overrides.values()): assert np.allclose(i.data, j.data, rtol=1.e-3) return op1 @switchconfig(ignore_unknowns=True) @pytest.mark.parametrize('extra,nqueues', [ ({}, 6), ({'deriv-uniformize': True}, 6), ]) def test_eiso_stress_like(self, extra, nqueues): self.run_eiso_stress_like(extra, nqueues) def test_multiple_x_derivs(self): grid = Grid(shape=(40, 40, 40)) x, y, z = grid.dimensions f = Function(name="f", grid=grid, space_order=8) g = f.func(name='g') u = TimeFunction(name="u", grid=grid, space_order=8) eqn = Eq(u.forward, u + f.dx + g.dy + g.dx) op0 = Operator(eqn, opt=opt4({'par-tile': (32, 8)})) op1 = Operator(eqn, opt=opt4({'errctl': 'basic', 'cire-hoist': 0, # Legacy 'gpu-regrot': 'scatter', 'par-tile': (32, 4)})) # Check generated code # Expected one regrot queue, incorporating g.dy into the partial # evaluations of f.dx self._check_scattering(op1, nqueues=2) assert 'queue0[qq0 - 4] += r0' in str(op1) u1 = u.func(name='u1') for i in [f, g, u, u1]: init_chessboard(i) op0.apply(time=3) op1.apply(time=3, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) def test_mul_as_dx_base(self): grid = Grid(shape=(40, 40, 40)) x, y, z = grid.dimensions f = Function(name='f', grid=grid, space_order=8) g = f.func(name='g') w0 = f.func(name='w0') w1 = f.func(name='w1') u = TimeFunction(name="u", grid=grid, space_order=8) eqn = Eq(u.forward, u + (f*w0).dx - w0.dz + (g*w1).dx) op0 = Operator(eqn, opt=opt4({'par-tile': (32, 4)})) op1 = Operator(eqn, opt=opt4({'errctl': 'basic', 'cire-hoist': 0, # Legacy 'gpu-regrot': 'scatter', 'par-tile': (32, 4)})) # Check generated code self._check_scattering(op1, nqueues=2) u1 = u.func(name='u1') for i in [f, g, w0, w1, u, u1]: init_chessboard(i) op0.apply(time=3) op1.apply(time=3, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-5) @pytest.mark.parametrize('space_order', [4, 16]) @pytest.mark.parametrize('uniformize', [False, True]) def test_summand_priming(self, space_order, uniformize): grid = Grid(shape=(24, 8, 8), extent=(23, 7, 7)) x, y, _ = grid.dimensions f = Function(name='f', grid=grid, space_order=space_order, staggered=y) g = Function(name='g', grid=grid, space_order=space_order, staggered=x) u = Function(name='u', grid=grid, space_order=space_order, staggered=(x, y)) # The centered g.dx widens g's X extent beyond the staggered f.dx # stencil, so g.dy needs a longer priming loop than f.dx eqn = Eq(u, (f.dx + g.dy)*g.dx/2) op = Operator(eqn, opt=opt4({'gpu-regrot': 'scatter', 'cire-mode': '1pass', 'par-tile': (16, 8, 8), 'deriv-uniformize': uniformize})) self._check_scattering(op) # Unit slopes, including halos: f.dx=0, g.dx=g.dy=1, hence u=0.5 f.data_with_halo[:] = 0 nx, ny, _ = g.data_with_halo.shape g.data_with_halo[:] = (np.arange(nx)[:, None, None] + np.arange(ny)[None, :, None]) op.apply() assert np.allclose(u.data, .5, atol=1.e-5, rtol=1.e-5) @pytest.mark.parametrize('gpu_vect', [True, 'stream']) def test_scatter_queue_initialization(self, gpu_vect): grid = Grid(shape=(31, 17, 25)) f = Function(name='f', grid=grid, space_order=4) g = f.func(name='g') u = f.func(name='u') v = f.func(name='v') u1 = u.func(name='u1') v1 = v.func(name='v1') eqns = [Eq(u, f.dx + f.dy), Eq(v, g.dx + g.dy)] options = {'cire-mode': '1pass', 'par-tile': (16, 4, 8)} op0 = Operator(eqns, name='op0', opt=opt4(options)) op1 = Operator(eqns, name='op1', opt=opt4({ **options, 'gpu-regrot': 'scatter', 'gpu-vect': gpu_vect })) # Scatter queues accumulate stencil contributions and must start at zero kernel = op1._func_table['kernel0'].root bundles = [i for i in FindSymbols().visit(kernel) if i.is_Bundle] queues = [i for i in bundles if isinstance(i.c0, TBQueue)] assert len(queues) == 1 assert queues[0].initvalue == (ListInitializer(0, 0),) assert any(i.components == (u, v) for i in bundles) for i in (f, g): init_chessboard(i) op0.apply() op1.apply(u=u1, v=v1) assert np.allclose(u.data, u1.data, rtol=1.e-6, atol=1.e-5) assert np.allclose(v.data, v1.data, rtol=1.e-6, atol=1.e-5) def test_scatter_weighted_derivatives(self): self.run_scatter_weighted_derivatives() def run_scatter_weighted_derivatives(self, gpu_glb_mem=False): grid = Grid(shape=(17, 17, 25), extent=(16., 16., 24.)) u = TimeFunction(name='u', grid=grid, space_order=4) v = Function(name='v', grid=grid, space_order=4) u0 = u.func(name='u0') eq = Eq(u.forward, 2*u.dx - 3*v.dx + u.dy) op0 = Operator(eq, name='scatter_weighted_ref', opt=('advanced', {'gpu-opt': False})) op1 = Operator(eq, name='scatter_weighted', opt=opt4({'gpu-glb-mem': gpu_glb_mem, 'gpu-prefetch': True, 'gpu-regrot': 'scatter', 'deriv-collect': False, 'cire-hoist': False, 'cire-mode': '1pass', 'par-tile': (32, 4, 8)})) # Scatter must preserve the coefficients of the other x-derivatives self._check_scattering(op1, nqueues=2) rng = np.random.default_rng(9028) u.data_with_halo[:] = rng.normal(size=u.data_with_halo.shape) u0.data_with_halo[:] = u.data_with_halo v.data_with_halo[:] = rng.normal(size=v.data_with_halo.shape) op0.apply(time_M=2, u=u0) op1.apply(time_M=2) np.testing.assert_allclose(u.data, u0.data, rtol=1.e-5, atol=1.e-4) return op1 @pytest.mark.parametrize('expr', [ '(v.dx + v.dy)/c + w.dx + w.dy', '((a*v).dy - (b*v).dx)/c + (a*w).dy - (b*w).dz', ], ids=['with-main', 'without-main']) def test_nested_derivative_sums(self, expr): self.run_nested_derivative_sums(expr) def run_nested_derivative_sums(self, expr, gpu_glb_mem=False): grid = Grid(shape=(17, 17, 25), extent=(16., 16., 24.)) x, y, z = grid.dimensions u = TimeFunction(name='u', grid=grid, space_order=8) v = u.func(name='v') w = u.func(name='w') u0 = u.func(name='u0') a = Function(name='a', grid=grid, space_order=8) b = a.func(name='b') c = Function(name='c', dimensions=(y,), shape=(17,), space_order=8) eq = Eq(u.forward, u + eval(expr)) op0 = Operator(eq, name='nested_derivative_sums_ref', opt=('advanced', {'gpu-opt': False})) op1 = Operator(eq, name='nested_derivative_sums', opt=opt4({'gpu-regrot': 'scatter', 'gpu-prefetch': True, 'gpu-glb-mem': gpu_glb_mem, 'gpu-hoist': 2, 'cire-mode': '1pass', 'par-tile': (32, 8)})) self._check_scattering(op1) rng = np.random.default_rng(9028) for f in (u, v, w): f.data_with_halo[:] = rng.normal(size=f.data_with_halo.shape) u0.data_with_halo[:] = u.data_with_halo for f in (a, b, c): f.data_with_halo[:] = rng.uniform(.5, 1.5, size=f.data_with_halo.shape) op0.apply(time_M=2, u=u0) op1.apply(time_M=2) np.testing.assert_allclose(u.data, u0.data, rtol=1.e-5, atol=1.e-4) return op1 @pytest.mark.parametrize('dtype', [ np.float32, pytest.param(np.float16, marks=skipif('nocudalike')) ]) def test_scatter_vector_queues(self, dtype): grid = Grid(shape=(17, 17, 25), extent=(16., 16., 24.)) u = TimeFunction(name='u', grid=grid, space_order=4, dtype=dtype) v = TimeFunction(name='v', grid=grid, space_order=4, dtype=dtype) u0 = u.func(name='u0') v0 = v.func(name='v0') eqns = [Eq(u.forward, u.dx + u.dy), Eq(v.forward, v.dx + v.dy)] sym_opt = {'half-arith': dtype == np.float16} op0 = Operator(eqns, name='scatter_vector_ref', sym_opt=sym_opt, opt=('advanced', {'gpu-opt': False})) op1 = Operator(eqns, name='scatter_vector', sym_opt=sym_opt, opt=opt4({'gpu-regrot': 'scatter', 'gpu-vect': True, 'deriv-collect': False, 'cire-hoist': False, 'cire-mode': '1pass', 'par-tile': (32, 4, 8)})) kernel = op1._func_table['kernel0'].root queues = [f for f in FindSymbols().visit(kernel) if f.is_Bundle and isinstance(f.c0, TBQueue)] assert len(queues) == 1 assert queues[0].ncomp == 2 assert queues[0].c0.dtype == dtype assert queues[0].initvalue == (ListInitializer(0, 0),) rng = np.random.default_rng(9028) for f, f0 in ((u, u0), (v, v0)): f.data_with_halo[:] = rng.normal(size=f.data_with_halo.shape) f0.data_with_halo[:] = f.data_with_halo op0.apply(time_M=2, u=u0, v=v0) op1.apply(time_M=2) rtol, atol = (1.e-2, 1.e-2) if dtype == np.float16 else (1.e-5, 1.e-4) for f, f0 in ((u, u0), (v, v0)): np.testing.assert_allclose(f.data, f0.data, rtol=rtol, atol=atol) class TestMPIOverlapMode: def _check_mpi_sequence(self, op, section=0): iters = FindNodes(Iteration).visit(op) time_iter, = [i for i in iters if i.dim.is_Time] section = time_iter.nodes[section] timed_list = section.body[0] body = timed_list.body[0].body assert isinstance(body[0], HaloUpdateList) assert body[1].is_Call # Compute assert isinstance(body[2], HaloWaitList) assert body[2].body[1].is_Call # sync qid1 assert isinstance(body[3].body[0], RemainderCall) assert body[3].body[1].is_Call # sync qid0 assert body[3].body[2].is_Call # sync qid1 @pytest.mark.parallel(mode=[(1, 'overlap2')]) def test_simple(self, mode): grid = Grid(shape=(30, 30, 30)) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=8) eqns = [ Eq(u.forward, u.dy + u.backward) ] op = Operator(eqns) op.cfunction # Check generated code self._check_mpi_sequence(op) @pytest.mark.parallel(mode=[(1, 'overlap2')]) def test_one_kernel_two_eqns_simple(self, mode): grid = Grid(shape=(30, 30, 30)) f = Function(name='f', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=8) eqns = [ Eq(f, u.dy), Eq(u.forward, u + u.backward + f) ] op = Operator(eqns) # Check generated code self._check_mpi_sequence(op) op.cfunction @pytest.mark.parallel(mode=[(1, 'overlap2')]) def test_two_kernels_simple(self, mode): grid = Grid(shape=(30, 30, 30)) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=8) v = u.func(name='v') eqns = [ Eq(u.forward, v.laplace + u + 1), Eq(v.forward, u.forward.laplace + v + 1) ] op = Operator(eqns) # Check generated code self._check_mpi_sequence(op) op.cfunction @pytest.mark.parallel(mode=[(4, 'overlap2')]) def test_too_small_to_run(self, mode): grid = Grid(shape=(31, 64, 25), topology=(1, 1, '*')) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=4) eqns = [ Eq(u.forward, u.dz + u.backward) ] op = Operator(eqns, opt=opt4({'errctl': 'basic', 'gpu-fit': [u]})) # MPI comp/overlap won'r work if the CORE region is too small because we # would launch empty kernels, which CUDA/HIP forbid with pytest.raises(ValueError): op.apply(time=10) @pytest.mark.parallel(mode=[(2, 'overlap2')]) @pytest.mark.parametrize('nx', [9, 11]) @pytest.mark.parametrize('compression', [None, 'noop']) def test_layered_domain_check(self, mode, nx, compression): grid = Grid(shape=(nx, 16), topology=(2, 1)) u = TimeFunction(name='u', grid=grid, space_order=0) usave = TimeFunction(name='usave', grid=grid, save=5, space_order=2, compression=compression, layers=DiskHost) op = Operator([Eq(u.forward, u + 1), Eq(usave, u)], opt=('advanced', {'buf-async-degree': 4}), name='op') assert not usave._is_finalized if nx == 9: # Only rank 1 has an empty core: both ranks must reject the run with pytest.raises(ValueError, match='total halo size'): op.apply(time_M=3) else: op.apply(time_M=3) assert np.all(u.data._local[0] == 4) assert np.all(u.data._local[1] == 3) @pytest.mark.parallel(mode=[(1, 'overlap2')]) def test_cire_kernel_w_custom_partile(self, mode): grid = Grid(shape=(30, 30, 30)) f = Function(name='f', grid=grid, space_order=8) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=8) v = u.func(name='v') eqns = [ Eq(u.forward, v.laplace + u + 1), Eq(v.forward, u.forward.laplace + f.dx + v + 1) ] par_tile = ((16, 8), (128, 4, 128), (64, 8, 96)) op = Operator(eqns, opt=opt4({'par-tile': par_tile})) # Check generated code self._check_mpi_sequence(op, section=0) self._check_mpi_sequence(op, section=1) op.cfunction @pytest.mark.parallel(mode=[(4, 'overlap2')]) @switchconfig(safe_math=True) # As nvc breaks with host_kwargs due to -fast def test_simple_stencil(self, host_kwargs, mode): grid = Grid(shape=(37, 45)) t = grid.stepping_dim x, y = grid.dimensions halo = ((0, 0), (2, 2), (2, 2)) u = TimeFunction(name='u', grid=grid, space_order=0, halo=halo) u1 = TimeFunction(name='u', grid=grid, space_order=0, halo=halo) eq = Eq(u.forward, u[t, x-1, y] + u[t, x, y] + u[t, x+1, y] + 1) op0 = Operator(eq, **host_kwargs) op1 = Operator(eq) # Check generated code assert all(i in op1._func_table for i in ['haloupdate0', 'halowait0']) op0.apply(time_M=9) op1.apply(time_M=9, u=u1) assert np.all(u.data == u1.data) @pytest.mark.parallel(mode=[(4, 'overlap2')]) @switchconfig(safe_math=True) # As nvc breaks with host_kwargs due to -fast def test_multiple_stencils(self, host_kwargs, mode): grid = Grid(shape=(18, 18)) t = grid.stepping_dim x, y = grid.dimensions u = TimeFunction(name='u', grid=grid, space_order=2) u1 = TimeFunction(name='u', grid=grid, space_order=2) x_stencil = u[t, x-2, y] + u[t, x-1, y] + u[t, x+1, y] + u[t, x+2, y] y_stencil = u[t, x, y-2] + u[t, x, y-1] + u[t, x, y+1] + u[t, x, y+2] eq = Eq(u.forward, u[t, x, y] + x_stencil + y_stencil + 1) op0 = Operator(eq, **host_kwargs) op1 = Operator(eq) op0.apply(time_M=5) op1.apply(time_M=5, u=u1) assert np.all(u.data == u1.data) @pytest.mark.parallel(mode=[(1, 'overlap2'), (4, 'overlap2')]) @pytest.mark.parametrize('so', [8, 16]) @pytest.mark.parametrize('side', [70, 133]) def test_alignment_core_and_remainder(self, mode, so, side): grid = Grid(shape=(side, side, side)) u = TimeFunction(name='u', grid=grid, time_order=2, space_order=so) eqn = Eq(u.forward, u.dz + 1) op = Operator(eqn, opt=opt4({'errctl': 'basic', 'gpu-vect': 'stream-minimal', 'par-tile': (32, 4, 64)})) # Check generated code self._check_mpi_sequence(op, section=0) # The Operator generates vector LOADs (i.e. float4-based LOADs) which # expect 16-byte alignment. The CORE and OWNED regions generated for # MPI computation/communication overlap must therefore start at a point # that is such that the vector LOADs honour the alignment contract. In # Devito, this is guaranteed by the `HaloScheme.omapper` lowering, which # takes the alignment requirements into account and optionally rounds # down the starting point of the CORE/OWNED regions to ensure alignment. # *Without* such a tweak, for example (`side==50` above), the remainder # regions would start/end at: # start/end at: # * z_m=0 z_M=15 # * z_m=34 z_M=49 # While with the tweak in place, they start/end at: # * z_m=0 z_M=15 # * z_m=32 z_M=47 op.apply(time=10, dt=.001) @pytest.mark.parametrize('exprs', [ (['u.dz2 + u.dy.dz', 'v.dz2 + v.dy.dz']), (['u.dx2 + u.dy2 + u.dz2 + u.dx.dy + u.dy.dz + u.dx.dz', 'v.dx2 + v.dy2 + v.dz2 + v.dx.dy + v.dy.dz + v.dx.dz']), ]) @pytest.mark.parametrize('topology', [ (1, '*', 1), (1, 1, '*'), ]) @pytest.mark.parametrize('cire_mode', [ '1pass', '2pass', ]) @pytest.mark.parametrize('gpu_vect', [ False, 'stream-minimal', ]) @pytest.mark.parallel(mode=[(4, 'overlap2')]) @switchconfig(safe_math=True) # As nvc breaks with host_kwargs due to -fast def test_mixed_derivs_incremental(self, host_kwargs, exprs, topology, cire_mode, gpu_vect, mode): grid = Grid(shape=(31, 64, 45), topology=topology) nt = 8 u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) v = TimeFunction(name='v', grid=grid, space_order=4, save=nt, layers=NoLayers) eqns = [] for expr, field in zip(as_tuple(exprs), [u, v]): eqns.append(Eq(field.forward, eval(expr) + 1.)) op0 = Operator(eqns, **host_kwargs) op1 = Operator(eqns, opt=opt4({'errctl': 'basic', 'par-tile': (16, 4), 'gpu-fit': [u, v], 'cire-mode': cire_mode, 'gpu-vect': gpu_vect})) # Check generated code assert 'compute0' in op1._func_table assert 'remainder0' in op1._func_table init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-2) @pytest.mark.parallel(mode=[(4, 'overlap2')]) @pytest.mark.parametrize('topology', [('*', '*', '*'), (1, 1, '*')]) def test_elastic_like(self, mode, topology): grid = Grid(shape=(100, 100, 100), topology=topology) v = VectorTimeFunction(name='v', grid=grid, space_order=4) tau = TensorTimeFunction(name='tau', grid=grid, space_order=4) eqns = [Eq(v.forward, v - div(tau)), Eq(tau.forward, tau - diag(div(v.forward)))] op = Operator(eqns, opt=opt4({'errctl': 'basic', 'gpu-vect': 'stream-minimal', 'par-tile': (32, 4, 64)})) # Check generated code self._check_mpi_sequence(op, section=0) self._check_mpi_sequence(op, section=1) # Ensure it compiles and runs through completion without any MPI or # data-alignment related issues op.apply(time=10, dt=.001) @pytest.mark.parametrize('expr,expected', [ ('u.dx2', 2), # stencils in x; topology=x; 2 comms expected ('u.laplace', 2), # stencils in x/y/z; topology=x; 2 comms expected ('u.dz2 + u.dy.dz', 0) # stencils in y/z; topology=x; no comms expected ]) @pytest.mark.parallel(mode=[(4, 'overlap2')]) @switchconfig(safe_math=True) # As nvc breaks with host_kwargs due to -fast def test_comms_based_on_topology(self, host_kwargs, expr, expected, mode): grid = Grid(shape=(70, 34, 35), topology=('*', 1, 1)) nt = 6 u = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) u1 = TimeFunction(name='u', grid=grid, space_order=4, save=nt, layers=NoLayers) eqn = Eq(u.forward, eval(expr) + 1.) op0 = Operator(eqn, **host_kwargs) op1 = Operator(eqn, opt=opt4({'errctl': 'basic', 'par-tile': (16, 4), 'gpu-fit': [u]})) # Check generated code hucs = FindNodes(HaloUpdateCall).visit(op1) if expected > 0: halo_update_call, = hucs assert halo_update_call.arguments[3] == expected remainder = op1._func_table['remainder0'].root comms_iter, = FindNodes(Iteration).visit(remainder) assert comms_iter.size() == expected else: assert not hucs init_chessboard(u) init_chessboard(u1) op0.apply(time_M=nt-2) op1.apply(time_M=nt-2, u=u1) assert np.allclose(u.data, u1.data, rtol=1.e-2) @pytest.mark.parallel(mode=[(2, 'overlap2')]) def test_sparse(self, mode): nt = 10 grid = Grid(shape=(10, 10)) f = TimeFunction(name='f', grid=grid) source = SparseTimeFunction(name='src', npoint=1, coordinates=[(0.5, 0.5)], nt=nt, grid=grid) src_term = source.inject(field=f.forward, expr=source) op = Operator([Eq(f.forward, f + 1)] + src_term) op(time=nt - 1, dt=1) @skipif('nocuda') class TestGraph: _specs = tuple(i for i in product(('left', 'middle', 'right'), repeat=3) if i != ('middle', 'middle', 'middle'))[:10] @staticmethod def _make_options(gpu_graph=True): return ('advanced', {'gpu-opt': True, 'gpu-graph': gpu_graph}) @classmethod def _make_equations(cls): grid = Grid(shape=(24, 24, 24)) fields = [] equations = [] for n, spec in enumerate(cls._specs): name = '_'.join(spec) subdomain = SpeccedDomain( f'boundary_{name}', spec, 4, grid=grid ) field = TimeFunction(name=f'u{n}', grid=subdomain, time_order=1) fields.append(field) equations.append(Eq(field.forward, field + n + 1)) return fields, equations def test_from_subdims(self): _, equations = self._make_equations() op = Operator(equations, name='graph_from_subdims', opt=self._make_options()) assert not FindNodes(KernelLaunch).visit(op) helpers = get_cuda_graphs(op) assert len(helpers) == 1 assert not FindNodes(CudaGraphKernelCall).visit(op) assert len([i for i in FindNodes(Call).visit(op) if i.name == helpers[0].name]) == 2 graph_calls = FindNodes(CudaGraphKernelCall).visit(helpers) assert len(graph_calls) == 10 code = str(op) assert '#include "devito_cuda_graph.h"' in code assert 'std::vector graph_nodes0(10);' in code assert code.count('devito_cuda_graph_kernel') == 10 assert 'devito_cuda_graph_dependency' not in code assert code.count('cudaGraphLaunch') == 1 def test_from_subdims_par_tile_mapping(self): _, equations = self._make_equations() par_tile = tuple((16 if i % 2 == 0 else 32, 4, 8) for i in range(len(equations))) op = Operator(equations, name='graph_from_subdims_par_tile', opt=('advanced', { 'gpu-opt': True, 'gpu-graph': True, 'par-tile': par_tile, })) calls = FindNodes(CudaGraphKernelCall).visit(get_cuda_graphs(op)) blocks = [str(i.block) for i in calls] assert blocks == ['block0', 'block1']*5 assert 'dim3 block0(16,4,8);' in str(op) assert 'dim3 block1(32,4);' in str(op) def test_from_subdims_disabled(self): _, equations = self._make_equations() op_no_graph = Operator( equations, name='graph_disabled', opt=self._make_options(gpu_graph=False) ) assert len(FindNodes(KernelLaunch).visit(op_no_graph)) == 10 assert not get_cuda_graphs(op_no_graph) assert 'devito_cuda_graph.h' not in str(op_no_graph) @switchconfig(profiling='advanced2') @pytest.mark.parametrize('fission', [False, 'basic']) def test_direct(self, fission): fields, equations = self._make_equations() op = Operator(equations, name='graph_direct', opt=opt0({ 'gpu-graph': 'direct', 'fission': fission, 'par-tile': ((32, 4, 8), (16, 4, 8))*5, })) assert len(FindNodes(KernelLaunch).visit(op)) == 10 assert not get_cuda_graphs(op) for horizon in (2, 0, 4): for field in fields: field.data[:] = 0 summary = op.apply(time_M=horizon) assert sum(len(v) for v in summary.subsections.values()) == 10 for n, field in enumerate(fields, 1): np.testing.assert_array_equal(field.data[0], horizon*n) np.testing.assert_array_equal(field.data[1], (horizon + 1)*n) @pytest.mark.parametrize('local', [False, True]) def test_subdim_locality(self, local): grid = Grid(shape=(8, 8, 8)) x, y, z = grid.dimensions zl = SubDimension.left('zl', z, 2, local=local) zr = SubDimension.right('zr', z, 2, local=local) a = TimeFunction(name='a', grid=grid) b = TimeFunction(name='b', grid=grid) eqns = [Eq(a.forward, a + 1).subs(z, zl), Eq(b.forward, b + 2).subs(z, zr)] op = Operator(eqns, name='graph_locality', opt=('advanced', { 'gpu-opt': True, 'gpu-graph': True, 'fission': 'basic' })) calls = FindNodes(CudaGraphKernelCall).visit(get_cuda_graphs(op)) assert len(calls) == (2 if local else 0) op.apply(time_M=0) assert np.all(a.data[1, :, :, :2] == 1) assert np.all(a.data[1, :, :, 2:] == 0) assert np.all(b.data[1, :, :, :-2] == 0) assert np.all(b.data[1, :, :, -2:] == 2) def test_from_subdims_execution(self): fields, equations = self._make_equations() op = Operator(equations, name='graph_execution', opt=self._make_options()) op.apply(time_M=2) for n, field in enumerate(fields, 1): actual = np.sort(field.data.reshape((2, -1)), axis=0) assert np.allclose(actual, ((2*n,), (3*n,))) # Graph handles belong to one invocation, not to the Operator lifetime for field in fields: field.data[:] = 0 op.apply(time_M=0) for n, field in enumerate(fields, 1): assert np.all(field.data[0] == 0) assert np.all(field.data[1] == n) @switchconfig(profiling='advanced') @pytest.mark.parametrize('backwards', [False, True]) @pytest.mark.parametrize('absolute_time', [False, True]) @pytest.mark.parametrize('time_m,time_M', [(0, 0), (2, 10), (5, 4)]) def test_setup(self, backwards, absolute_time, time_m, time_M): grid = Grid(shape=(17, 19, 23)) left = SpeccedDomain('left', ('left', 'none', 'none'), 4, grid=grid) right = SpeccedDomain('right', ('right', 'none', 'none'), 4, grid=grid) a = TimeFunction(name='a', grid=grid, time_order=1) b = TimeFunction(name='b', grid=grid, time_order=2) a1 = a.func(name='a1') b1 = b.func(name='b1') value = grid.time_dim if absolute_time else 1 previous = b.forward if backwards else b.backward eqs = [Eq(a.backward if backwards else a.forward, a + value, subdomain=left), Eq(b.backward if backwards else b.forward, b + previous + value, subdomain=right)] options = {'gpu-opt': True, 'fission': 'basic'} ref = Operator(eqs, name='graph_setup_ref', opt=('advanced', {**options, 'gpu-graph': False})) op = Operator(eqs, name='graph_setup', opt=('advanced', {**options, 'gpu-graph': True})) helper, = get_cuda_graphs(op) execs = [s for s in FindSymbols().visit(op) if isinstance(s, CudaGraphExec)] assert len(execs) == 1 # Build and upload once, then update kernel arguments before each launch iteration, = FindNodes(Iteration).visit(op) timed, = FindNodes(TimedList).visit(iteration) assert FindNodes(TimedList).visit(op) == [timed] assert [c.name for c in FindNodes(Call).visit(timed)] == [ helper.name, 'cudaGraphLaunch', 'cudaStreamSynchronize' ] assert 'cudaGraphUpload' not in [c.name for c in FindNodes(Call).visit(iteration)] calls = [c.name for c in FindNodes(Call).visit(op)] assert calls.count(helper.name) == 2 assert calls.count('cudaGraphUpload') == 1 # Repeat with another time range, rebuilding all captured arguments for start, stop in [(time_m, time_M), (3, 9)]: for f in (a, b, a1, b1): init_chessboard(f) ref.apply(time_m=start, time_M=stop) op.apply(time_m=start, time_M=stop, a=a1, b=b1) np.testing.assert_allclose(a1.data, a.data) np.testing.assert_allclose(b1.data, b.data) @pytest.mark.parametrize('backwards', [False, True]) def test_setup_sequential_dimension(self, backwards): grid = Grid(shape=(8, 8, 8)) left = SpeccedDomain('left', ('left', 'none', 'none'), 2, grid=grid) right = SpeccedDomain('right', ('right', 'none', 'none'), 2, grid=grid) i = Dimension(name='i') a = Function(name='a', grid=grid, dimensions=(i, *grid.dimensions), shape=(5, *grid.shape), space_order=0) b = a.func(name='b') next_i = i - 1 if backwards else i + 1 op = Operator([Eq(a.subs(i, next_i), a + 1, subdomain=left), Eq(b.subs(i, next_i), b + 2, subdomain=right)], name='graph_sequential_dimension', opt=self._make_options()) iteration, = FindNodes(Iteration).visit(op) assert iteration.dim is i assert iteration.is_Sequential assert 'cudaGraphUpload' in [c.name for c in FindNodes(Call).visit(op)] assert 'cudaGraphUpload' not in [c.name for c in FindNodes(Call).visit(iteration)] op.apply(i_m=1 if backwards else 0, i_M=4 if backwards else 3) expected = np.arange(5)[::-1] if backwards else np.arange(5) expected = expected[:, None, None, None] assert np.all(a.data[:, :2] == expected) assert np.all(a.data[:, 2:] == 0) assert np.all(b.data[:, -2:] == 2*expected) assert np.all(b.data[:, :-2] == 0) @pytest.mark.parametrize('gpu_graph', [False, True]) def test_dynamic_shm(self, gpu_graph): """Opt in to large dynamic shared memory for ordinary and graph launches.""" grid = Grid(shape=(17, 19, 65), extent=(16., 18., 64.)) domains = [SpeccedDomain(f'{side}_slab', ('none', 'none', side), 10, grid=grid) for side in ('left', 'right')] inputs = [TimeFunction(name=f'p{i}', grid=grid, space_order=8) for i in range(6)] fields = [TimeFunction(name=f'u{i}', grid=v, space_order=8) for i, v in enumerate(domains)] refs = [f.func(name=f'{f.name}_ref') for f in fields] stencil = sum(p.dx + p.dy + p.dz for p in inputs) eqs = [Eq(f.forward, f + stencil) for f in fields] op0 = Operator(eqs, name='graph_shm_ref', opt=opt0({'gpu-graph': False})) op1 = Operator(eqs, name='graph_shm_opt', opt=('advanced', { 'errctl': 'basic', 'gpu-opt': True, 'gpu-graph': gpu_graph, 'gpu-vect': 'stream-only', 'gpu-glb-mem': 'async-pipe', 'gpu-prefetch': 3, 'gpu-wdr': False, 'par-tile': (64, 4, 48), 'fission': False, })) assert str(op1).count('cudaFuncSetAttribute') == 2 for f in [*inputs, *fields, *refs]: init_chessboard(f) op0.apply(time_M=2) op1.apply(time_M=2, **{f.name: ref for f, ref in zip(fields, refs)}) for f, ref in zip(fields, refs): np.testing.assert_allclose(f.data, ref.data, rtol=1.e-5, atol=1.e-6) @pytest.mark.parallel(mode=[(4, 'basic')]) def test_from_empty_subdims(self, mode): fields, equations = self._make_equations() # MPI leaves one rank with an entirely empty graph, which must remain # valid and profile cleanly without overriding thicknesses op = Operator(equations, name='graph_empty', opt=self._make_options()) comm = fields[0].grid.comm # Allocate every field on every rank before testing for emptiness sizes = [i.data.size for i in fields] empty = all(i == 0 for i in sizes) assert comm.allreduce(int(empty)) == 1 op.apply(time_M=0) for n, field in enumerate(fields, 1): assert np.all(field.data[0]._local == 0) assert np.all(field.data[1]._local == n) @pytest.mark.parallel(mode=[(4, 'basic')]) def test_from_partially_empty_subdims_mpi(self, mode): grid = Grid(shape=(24, 24, 24), topology=(2, 2, 1)) x, y, _ = grid.dimensions xface = SpeccedDomain('xface', ('left', 'none', 'none'), 4, grid=grid) yface = SpeccedDomain('yface', ('none', 'left', 'none'), 4, grid=grid) def make_equations(prefix): a = TimeFunction(name=f'{prefix}a', grid=grid, time_order=1) b = TimeFunction(name=f'{prefix}b', grid=grid, time_order=1) c = TimeFunction(name=f'{prefix}c', grid=grid, time_order=1) equations = [ Eq(a.forward, a + 1, subdomain=xface), Eq(b.forward, a.forward + 1, subdomain=yface), Eq(c.forward, a.forward + b.forward, subdomain=xface) ] return (a, b, c), equations expected, reference_equations = make_equations('ref_') actual, graph_equations = make_equations('graph_') reference = Operator( reference_equations, name='graph_partial_empty_ref', opt=self._make_options(gpu_graph=False) ) op = Operator(graph_equations, name='graph_partial_empty', opt=self._make_options()) helpers = get_cuda_graphs(op) assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 3 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] assert len(dependencies) == 3 # Exercise a rank where the middle node is empty but its predecessor # and successor are not positions = grid.distributor.glb_pos_map bypasses_middle = LEFT in positions[x] and LEFT not in positions[y] assert grid.distributor.comm.allreduce(int(bypasses_middle)) == 1 reference.apply(time_M=2) op.apply(time_M=2) # Compare rank-local NumPy views; every MPI rank runs the assertion for expected_field, actual_field in zip(expected, actual, strict=True): actual_data = np.asarray(actual_field.data._local) expected_data = np.asarray(expected_field.data._local) np.testing.assert_allclose(actual_data, expected_data) @pytest.mark.parametrize('shared_boundary', [False, True]) @pytest.mark.parametrize('shared_output', [False, True]) def test_from_disjoint_subdims(self, shared_boundary, shared_output): grid = Grid(shape=(24, 24, 24)) x, y, z = grid.dimensions xl = SubDimension.left('xl', x, 4) yl = SubDimension.left('yl', y, 4) zl = SubDimension.left('zl', z, 4, separated=True) if shared_boundary: zs = SubDimension.middle('zm', z, 4, 4, separated=True) zs = zs._rebuild(thickness=(zl.ltkn, zs.rtkn)) else: zs = SubDimension.right('zr', z, 4, separated=True) a = TimeFunction(name='a', grid=grid, time_order=1) b = TimeFunction(name='b', grid=grid, time_order=1) output = a if shared_output else b # Sharing `a` does not induce a dependence across disjoint regions disjoint = Operator([ Eq(a.forward, a + 1).subs({x: xl, y: yl, z: zl}), Eq(output.forward, a.forward + 1).subs({x: xl, y: yl, z: zs}) ], name='graph_from_disjoint_subdims', opt=self._make_options()) assert not FindNodes(KernelLaunch).visit(disjoint) helpers = get_cuda_graphs(disjoint) assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 2 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] assert not dependencies assert '#include "devito_cuda_graph.h"' in str(disjoint) disjoint.apply(time_M=0) assert np.max(a.data) == 1 assert np.max(b.data) == (0 if shared_output else 1) @pytest.mark.parametrize('side', ['left', 'right']) def test_from_disjoint_faces(self, side): grid = Grid(shape=(32, 32)) xface = SpeccedDomain('xface', (side, 'middle'), 4, grid=grid) yface = SpeccedDomain('yface', ('middle', side), 4, grid=grid) u = TimeFunction(name='u', grid=grid, time_order=1, space_order=0) eqns = [Eq(u.forward, 1, subdomain=xface), Eq(u.forward, 2, subdomain=yface)] op = Operator(eqns, name='graph_disjoint_faces', opt=self._make_options()) # Distinct thickness symbols describe disjoint faces of one partition helpers = get_cuda_graphs(op) assert len(helpers) == 1 assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 2 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] assert not dependencies op.apply(time_M=0) expected = np.zeros(grid.shape) slab = slice(None, 4) if side == 'left' else slice(-4, None) expected[slab, 4:-4] = 1 expected[4:-4, slab] = 2 slices = tuple(grid.distributor.glb_slices[d] for d in grid.dimensions) np.testing.assert_allclose(u.data[1]._local, expected[slices]) @pytest.mark.parametrize('side', ['left', 'right']) @pytest.mark.parallel(mode=[(4, 'basic')]) def test_from_disjoint_faces_mpi(self, side, mode): self.test_from_disjoint_faces(side) @pytest.mark.parametrize('update', [Eq, Inc], ids=['write', 'increment']) def test_from_overlapping_subdims(self, update): grid = Grid(shape=(24, 24, 24)) xface = SpeccedDomain('xface', ('left', 'none', 'none'), 4, grid=grid) yface = SpeccedDomain('yface', ('none', 'left', 'none'), 4, grid=grid) a = TimeFunction(name='a', grid=grid, time_order=1) op = Operator([ update(a.forward, 1, subdomain=xface), update(a.forward, 2, subdomain=yface) ], name=f'graph_overlapping_{update.__name__.lower()}', opt=self._make_options()) helpers = get_cuda_graphs(op) assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 2 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] assert len(dependencies) == 1 op.apply(time_M=0) expected = np.zeros_like(a.data[1]) expected[:4, :, :] = 1 if update is Eq: expected[:, :4, :] = 2 else: expected[:, :4, :] += 2 assert np.all(a.data[1] == expected) def test_time_aliases(self): grid = Grid(shape=(24, 24, 24)) xface = SpeccedDomain('xface', ('left', 'none', 'none'), 4, grid=grid) yface = SpeccedDomain('yface', ('none', 'left', 'none'), 4, grid=grid) a = TimeFunction(name='a', grid=grid, time_order=2, save=Buffer(2)) # Forward and backward address the same slot in a two-slot cyclic buffer op = Operator([ Eq(a.forward, 1, subdomain=xface), Eq(a.backward, 2, subdomain=yface) ], name='graph_time_aliases', opt=self._make_options()) helpers = get_cuda_graphs(op) assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 2 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] assert len(dependencies) == 1 op.apply(time_m=1, time_M=1) expected = np.zeros_like(a.data) expected[0, :4, :, :] = 1 expected[0, :, :4, :] = 2 np.testing.assert_allclose(a.data, expected) @pytest.mark.parametrize('override', ['function', 'bounds', 'thickness', 'empty']) @pytest.mark.parametrize('shift', [0, -1]) def test_subdomain_runtime_overrides(self, override, shift): grid = Grid(shape=(8, 8, 32)) z = grid.dimensions[-1] zl = SubDimension.left('zl', z, 8, separated=True) zr = SubDimension.right('zr', z, 8, separated=True) a = TimeFunction(name='a', grid=grid) b = TimeFunction(name='b', grid=grid) op = Operator([ Eq(a.forward, a + 1).subs(z, zl), Eq(b.forward, a.forward.subs(z, z + shift) + 1).subs(z, zr) ], name='graph_runtime_overrides', opt=self._make_options()) helpers = get_cuda_graphs(op) assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 2 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] # The runtime contract protects inward reads within space_order assert len(dependencies) == 0 kwargs = {} left = 8 m, M = 0, 31 if override == 'function': runtime_grid = Grid(shape=(8, 8, 17)) a = TimeFunction(name='runtime_a', grid=runtime_grid) b = TimeFunction(name='runtime_b', grid=runtime_grid) kwargs.update(a=a, b=b) M = 16 elif override == 'bounds': with pytest.raises(InvalidArgument, match='Cannot override bounds'): op.apply(time_M=0, z_m=3, z_M=19) return else: left = 23 if override == 'thickness' else 0 tkn, = [i for i in op.parameters if isinstance(i, Thickness) and i.side is LEFT] kwargs[tkn.name] = left with pytest.raises(InvalidArgument, match='Cannot override SubDimension thickness'): op.apply(time_M=0, **kwargs) return op.apply(time_M=0, **kwargs) expected_a = np.zeros_like(a.data[1]) expected_a[:, :, m:m + left] = 1 expected_b = np.zeros_like(b.data[1]) expected_b[:, :, M - 7:M + 1] = expected_a[:, :, M - 7 + shift:M + 1 + shift] + 1 np.testing.assert_allclose(a.data[1], expected_a) np.testing.assert_allclose(b.data[1], expected_b) @pytest.mark.parametrize('separated,thickness,shift', [ (True, 8, -1), (False, 10, 0), (None, 10, 0) ]) def test_left_right_contract(self, separated, thickness, shift): grid = Grid(shape=(8, 8, 16)) z = grid.dimensions[-1] zl = SubDimension.left('zl', z, thickness, separated=True) kwargs = {} if separated is None else {'separated': separated} zr = SubDimension.right('zr', z, thickness, **kwargs) a = TimeFunction(name='a', grid=grid) b = TimeFunction(name='b', grid=grid) eqs = [Eq(a.forward, 1).subs(z, zl), Eq(b.forward, a.forward.subs(z, z + shift) + 1).subs(z, zr)] op = Operator(eqs, name='graph_left_right_contract', opt=self._make_options()) helpers = get_cuda_graphs(op) assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 2 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] assert len(dependencies) == (0 if separated else 1) if separated: with pytest.raises(InvalidArgument, match='gap of at least 1'): op.apply(time_M=0) return op.apply(time_M=0) expected_a = np.zeros(grid.shape) expected_a[:, :, :thickness] = 1 expected_b = np.zeros(grid.shape) start = 16 - thickness expected_b[:, :, start:] = expected_a[:, :, start + shift:16 + shift] + 1 np.testing.assert_allclose(a.data[1], expected_a) np.testing.assert_allclose(b.data[1], expected_b) @pytest.mark.parametrize('separated', [False, True]) def test_stencil_gap(self, separated): grid = Grid(shape=(8, 8, 64)) z = grid.dimensions[-1] zl = SubDimension.left('zl', z, 10, separated=True) zr = SubDimension.right('zr', z, 18, separated=separated) u = TimeFunction(name='u', grid=grid, space_order=4) v = u.func(name='v') eqs = [Eq(u.forward, 1).subs(z, zl), Eq(v.forward, u.forward.subs(z, z - 4)).subs(z, zr)] op = Operator(eqs, name='graph_stencil_gap', opt=self._make_options()) helpers = get_cuda_graphs(op) assert len(FindNodes(CudaGraphKernelCall).visit(helpers)) == 2 dependencies = [i for i in FindNodes(Call).visit(helpers) if i.name == 'devito_cuda_graph_dependency'] assert len(dependencies) == (0 if separated else 1) size = 32 if separated else 28 runtime_grid = Grid(shape=(8, 8, size), dimensions=grid.dimensions) fields = [TimeFunction(name=f'{f.name}_runtime', grid=runtime_grid, space_order=4) for f in (u, v)] op.apply(time_M=0, **dict(zip(('u', 'v'), fields))) expected_u = np.zeros(runtime_grid.shape) expected_u[:, :, :10] = 1 expected_v = np.zeros(runtime_grid.shape) expected_v[:, :, -18:] = expected_u[:, :, size - 22:size - 4] slices = tuple(runtime_grid.distributor.glb_slices[d] for d in runtime_grid.dimensions) for field, expected in zip(fields, (expected_u, expected_v)): np.testing.assert_allclose(field.data[1]._local, expected[slices]) @pytest.mark.parallel(mode=[(4, 'basic')]) def test_stencil_gap_mpi(self, mode): self.test_stencil_gap(True) @skipif('nocuda') def test_par_tile_per_node(self): grid = Grid(shape=(16, 16, 64)) x, y, z = grid.dimensions ix = SubDimension.left('ix', x, 4) iy = SubDimension.middle('iy', y, 2, 2, local=True) iz = SubDimension.middle('iz', z, 2, 2, local=True) c0 = ConditionalDimension('c0', ix, condition=Lt(ix, 2)) c1 = ConditionalDimension('c1', ix, condition=Gt(ix, 1)) f = TimeFunction(name='f', grid=grid) g = TimeFunction(name='g', grid=grid) # Different guards keep two Clusters with the same blocked iteration space eqs = [Eq(f.forward, f + 1, implicit_dims=(c0,)), Eq(g.forward, g + 2, implicit_dims=(c1,))] eqs = [e.subs({x: ix, y: iy, z: iz}) for e in eqs] tiles = ((16, 4, 8), (32, 4, 8)) op = Operator(eqs, name='graph_par_tile_per_node', opt=opt0({ 'par-tile': tiles, 'gpu-graph': True, 'fission': False })) helpers = get_cuda_graphs(op) assert len(helpers) == 1 nodes = FindNodes(CudaGraphKernelCall).visit(helpers) assert len(nodes) == 2 assert tuple(i.arguments[5].cargs for i in nodes) == tiles op.apply(time_M=0) for field, region, value in ((f, slice(0, 2), 1), (g, slice(2, 4), 2)): expected = np.zeros_like(field.data) expected[1, region, 2:-2, 2:-2] = value np.testing.assert_array_equal(field.data, expected) class TestProfiling: def test_ops(self): """ Make sure indexing operations such as int x = (int)(blockIdx.x*blockDim.x + threadIdx.x); aren't counted as FLOPs. """ grid = Grid(shape=(4, 4)) u = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u + 1.) op = Operator(eq) section0 = op._profiler._sections['section0'] assert section0.sops == 1 assert len(section0.itermaps) == 1 @switchconfig(profiling='advanced') def test_itershapes(self, host_kwargs): shape = (80, 80, 80) grid = Grid(shape=shape) u = TimeFunction(name='u', grid=grid) eq = Eq(u.forward, u.dx + 1.) tb = (32, 4, 1) op = Operator(eq, opt=opt0({'par-tile': tb})) summary = op.apply(time_M=9) assert len(summary) == 1 _, v = summary.popitem() itershapes = v.itershapes assert len(itershapes) == 2 itershape = itershapes[0] assert len(itershape) == 8 # time, 3 outer block dims, 3 inner block dims assert itershape[1:4] == tuple(reversed([i // j for i, j in zip(shape, tb)])) assert itershape[4:7] == tuple(reversed(tb)) # Should be the same as on host (but flipped) with runtime args op1 = Operator(eq, opt=('advanced', {'blockinner': True}), **host_kwargs) summary1 = op1.apply(time_M=9, x0_blk0_size=tb[0], y0_blk0_size=tb[1], z0_blk0_size=tb[2]) _, v1 = summary1.popitem() assert v1.itershapes[0] == ((itershape[0],) + tuple(reversed(itershape[1:4])) + tuple(reversed(itershape[4:7]))) @switchconfig(profiling='advanced') def test_gpoints_gflops_kernopt_default(self, host_kwargs): shape = (80, 80, 80) grid = Grid(shape=shape) u = TimeFunction(name='u', grid=grid, space_order=2) eq0 = Eq(u.forward, u.dx + 1.) eq1 = Eq(u.forward, u.dx + u.dy2 + 1.) tb = (32, 4) op0 = Operator(eq0, opt=opt4({'par-tile': tb})) op1 = Operator(eq1, opt=opt4({'par-tile': tb})) section00 = op0._profiler._sections['section0'] section01 = op1._profiler._sections['section0'] args = normalize_args(op0.arguments(time_M=0)) # Expected the same number of grid points, of course! assert section00.points.subs(args) == np.prod(grid.shape) assert section00.points == section01.points # Now test the gflops -- should be the same as on host! op2 = Operator(eq1, **host_kwargs) section02 = op2._profiler._sections['section0'] # +1 because of a different factorization strategy in PRO assert section01.sops == section02.sops + 1 assert section01.points.subs(args) == section02.points.subs(args) @pytest.mark.skip("Vector types not working at the moment") @switchconfig(profiling='advanced') def test_gpoints_gflops_kernopt_vecttypes(self, host_kwargs): shape = (80, 80, 80) grid = Grid(shape=shape) u = TimeFunction(name='u', grid=grid, space_order=2) v = TimeFunction(name='v', grid=grid, space_order=2) eq = Eq(u.forward, u.dx + v.dx + 1.) # W/ and w/o vector types should yield the same points! tb = (32, 4) op0 = Operator(eq, opt=opt0({'par-tile': tb})) op1 = Operator(eq, opt=opt5({'par-tile': tb})) section00 = op0._profiler._sections['section0'] section01 = op1._profiler._sections['section0'] args = normalize_args(op0.arguments(time_M=0)) # Expected the same number of grid points and operations assert section00.points.subs(args) == np.prod(grid.shape) assert section00.points == section01.points assert section00.ops == section01.ops # Now test the gflops -- should be the same as on host! op2 = Operator(eq, **host_kwargs) section02 = op2._profiler._sections['section0'] assert section01.sops == section02.sops assert section01.points.subs(args) == section02.points.subs(args) assert section01.ops.subs(args) == section02.ops.subs(args) # TODO: unfortunately at the moment the traffic may be overestimated # in the case of vector types if one writes to a subset of components # only. However, this is not a typical case in real-life scenarios assert section01.traffic.subs(args) > section02.traffic.subs(args) if __name__ == "__main__": configuration['mpi'] = 'diag2' _host_kwargs = {'platform': 'intel64', 'compiler': 'nvc', 'language': 'C'} TestKernelOptVectTypes().test_carry_over_halo_metadata()