compiler: Patch factorizer pass after relaxing is_dense

devitocodes · Aug 2, 2022 · 693485f · 693485f
1 parent f8380b6
commit 693485f
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/.github/workflows/docker-devito.yml b/.github/workflows/docker-devito.yml
@@ -22,13 +22,13 @@ jobs:
             tag: 'nvidia-nvc'
             flag: '--gpus all'
             test: 'tests/test_gpu_openacc.py tests/test_gpu_common.py'
-            runner: ["self-hosted", "gpu", "docker"]
+            runner: ["self-hosted", "gpu", "docker", "v1004"]
 
           - base: 'bases:nvidia-clang'
             tag: 'nvidia-clang'
             flag: '--gpus all'
             test: 'tests/test_gpu_openmp.py tests/test_gpu_common.py'
-            runner: ["self-hosted", "gpu", "docker"]
+            runner: ["self-hosted", "gpu", "docker", "kimogila"]
 
           # Runtime gpu flags from https://hub.docker.com/r/rocm/tensorflow/
           - base: 'bases:amd'

diff --git a/devito/passes/clusters/factorization.py b/devito/passes/clusters/factorization.py
@@ -3,7 +3,7 @@
 from sympy import Add, Mul, S, collect
 
 from devito.ir import cluster_pass
-from devito.symbolics import estimate_cost, retrieve_symbols
+from devito.symbolics import BasicWrapperMixin, estimate_cost, retrieve_symbols
 from devito.tools import ReducerMap
 
 __all__ = ['factorize']
@@ -115,7 +115,7 @@ def run(expr):
             return expr, {'funcs': expr}
         elif expr.is_Pow:
             return expr, {'pows': expr}
-        elif expr.is_Symbol or expr.is_Indexed or expr.is_Atom:
+        elif expr.is_Symbol or expr.is_Indexed or isinstance(expr, BasicWrapperMixin):
             return expr, {}
         elif expr.is_Add:
             args, candidates = zip(*[run(arg) for arg in expr.args])

diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py
@@ -15,7 +15,8 @@
            'FieldFromComposite', 'ListInitializer', 'Byref', 'IndexedPointer', 'Cast',
            'DefFunction', 'InlineIf', 'Keyword', 'String', 'Macro', 'MacroArgument',
            'CustomType', 'Deref', 'INT', 'FLOAT', 'DOUBLE', 'VOID', 'CEIL',
-           'FLOOR', 'MAX', 'MIN', 'Null', 'SizeOf', 'rfunc', 'cast_mapper']
+           'FLOOR', 'MAX', 'MIN', 'Null', 'SizeOf', 'rfunc', 'cast_mapper',
+           'BasicWrapperMixin']
 
 
 class CondEq(sympy.Eq):