diff --git a/devito/passes/clusters/cse.py b/devito/passes/clusters/cse.py index 87770d4039..9c899eb23e 100644 --- a/devito/passes/clusters/cse.py +++ b/devito/passes/clusters/cse.py @@ -14,7 +14,7 @@ from devito.finite_differences.differentiable import IndexSum from devito.ir import Cluster, Scope, cluster_pass from devito.symbolics import ( - DefFunction, Reserved, estimate_cost, q_leaf, q_terminal, search + DefFunction, IndexedPointer, Reserved, estimate_cost, q_leaf, q_terminal, search ) from devito.symbolics.manipulation import _uxreplace from devito.tools import DAG, as_list, as_tuple, extract_dtype, frozendict @@ -418,11 +418,17 @@ def _(expr): @_catch.register(Indexed) +@_catch.register(IndexedPointer) @_catch.register(Symbol) @_catch.register(Reserved) def _(expr): """ Handler for objects preventing CSE to propagate through their arguments. + + For IndexedPointer (e.g. `f_vec->sizes[(time / ct0f) + ...]`) this also + keeps integer index expressions out of a floating-point cluster's CSE + temporaries, which `cse_dtype` would otherwise declare as floating-point + and thus produce non-integral array subscripts. """ return {} diff --git a/tests/test_cse.py b/tests/test_cse.py index 976593e9ff..6ab68f05a0 100644 --- a/tests/test_cse.py +++ b/tests/test_cse.py @@ -11,7 +11,7 @@ from devito.ir import Conditional, DummyEq, FindNodes, FindSymbols from devito.ir.support import generator from devito.passes.clusters.cse import CTemp, _cse -from devito.symbolics import DefFunction, indexify +from devito.symbolics import DefFunction, FieldFromPointer, IndexedPointer, indexify from devito.types import Array, Symbol, Temp @@ -154,6 +154,33 @@ def test_deffunction_not_captured(): assert processed == exprs +def test_indexedpointer_index_not_captured(): + """ + The index of an IndexedPointer (e.g. `f_vec->sizes[...]`) must not be + CSE-captured: in a floating-point cluster `cse_dtype` would bind the + integer index to a floating-point temporary, yielding a non-integral + array subscript that does not compile (nvcc, hipcc). + """ + time = Symbol(name='time', dtype=np.int32) + fac = Symbol(name='fac', dtype=np.int32) + vec = Symbol(name='f_vec', dtype=np.int32) + a = CTemp(name='a', dtype=np.float32) + b = CTemp(name='b', dtype=np.float32) + + index = time/fac + 2*time + sizes = IndexedPointer(FieldFromPointer('sizes', vec), index) + offs = IndexedPointer(FieldFromPointer('offs', vec), index) + exprs = [DummyEq(a, sizes + 1.), DummyEq(b, offs + sizes)] + + counter = generator() + make = lambda e: CTemp(name=f'r{counter()}', dtype=np.float32) + processed = _cse(exprs, make) + + for e in processed: + for i in e.find(IndexedPointer): + assert not any(isinstance(s, CTemp) for s in i.index.free_symbols) + + def test_w_conditionals(): grid = Grid(shape=(10, 10, 10)) x, _, _ = grid.dimensions