compiler: Revamp lowering of IndexDerivatives #2310

FabioLuporini · 2024-02-12T14:08:55Z

This paves the way for external customisation, e.g.

w[i]*u[x + i] vs w[i]*(u[x + i] - u[x - i])
loops vs unroll
etc

FabioLuporini · 2024-02-12T14:10:08Z

tests/test_unexpansion.py

@@ -324,10 +324,12 @@ def test_redundant_derivatives(self):
        assert len(get_arrays(op)) == 0
        assert op._profiler._sections['section0'].sops == 74
        exprs = FindNodes(Expression).visit(op)
-        assert len(exprs) == 6
+        assert len(exprs) == 5


Note for reviewers:

this is a result of this tweak: https://github.com/devitocodes/devito/pull/2310/files#diff-a012e847c0c9841809763e4c55eb89ee0ee9b985614693494911e4ab32bda7f1R193-R194

codecov · 2024-02-12T14:13:10Z

Codecov Report

Attention: 8 lines in your changes are missing coverage. Please review.

Comparison is base (c755cd0) 86.67% compared to head (8f22f90) 86.67%.
Report is 2 commits behind head on master.

Files	Patch %	Lines
devito/core/operator.py	42.85%	2 Missing and 2 partials ⚠️
devito/passes/clusters/derivatives.py	93.10%	3 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #2310   +/-   ##
=======================================
  Coverage   86.67%   86.67%           
=======================================
  Files         229      229           
  Lines       43107    43150   +43     
  Branches     7996     8001    +5     
=======================================
+ Hits        37363    37401   +38     
- Misses       5050     5053    +3     
- Partials      694      696    +2

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

mloubout · 2024-02-12T14:32:59Z

devito/finite_differences/differentiable.py

@@ -803,9 +815,13 @@ def _evaluate(self, **kwargs):
        return EvalDerivative(expr, base=self.base)


+class DiffDerivative(IndexDerivative, DifferentiableOp):


Why not make IndexDerivative directlu DifferentiableOp

I tried all sort of ways; remember we discussed a sympy issue a few days ago, and eventually I landed here...

So, IndexDerivative was (master branch) a DifferentiableOp, since IndexSum was (master branch) a DifferentiableOp. The heart of the problem is: if an object is non differentiable but one of its arguments is, then a reconstruction will make the parent object differentiable as well

This new structuring distinguishes between DiffDerivative (differentiable) and the lower-level IndexDerivative (non-differentiable). This ensures that, once inside the compiler, everything will have been lowered such that no Differentiable objects can ever be (re)constructed, which in essence will make compilation behave as expected, having ruled out by construction any funky-ness due to the way SymPy handles subclasses and reconstruction

Does it make sense or does it sound like a ranting? I can elaborate or explain in a call if necessary

mloubout · 2024-02-12T14:33:47Z

devito/finite_differences/differentiable.py

+
+    @property
+    def scaling(self):
+        return Mul(*[d.spacing**self.order for d in self.mapper])


While yes we wanna go there, this doesn't belong here, this belongs in Derivatives so we can progressively make the eval_fd dimension-indep fir easier compiler pass irrespective of expand/no-expand

This expression will also be wrong for an IndexDerivative along a diagonal such as

https://github.com/devitocodespro/recipes/blob/d422c271d167cf6a516fa0f8ae7c323ba9b74963/recipes/elastic_tti/fd_utils.py#L61

I think I can drop it, will check

mloubout · 2024-02-12T14:35:59Z

devito/finite_differences/finite_difference.py

@@ -247,7 +247,7 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, symbolic
            # Pure number
            pass

-        deriv = IndexDerivative(expr*weights, {dim: indices.free_dim})
+        deriv = DiffDerivative(expr*weights, {dim: indices.free_dim}, deriv_order)


I think we should drop the if exand/else since can just do this part and return deriv.evaluate if expand=true since IndexDerivative are properly evaluable since a few PR ago

I worry it opens a Pandora's box that I wouldn't be very willing to deal with at this stage, but I can try

I'm now opening a different PR for this

How do u suggest we handle the second part of the if:

if not expand and indices.expr is not None:

IIRC indices.expr may be None for some low order derivatives.
I think we would first have to normalizate behavior at the top (always produce indices.expr) then we could do what you suggested

Hum right, didn't think that indices.expr would be None I guess we can park it for now

mloubout · 2024-02-12T14:36:20Z

devito/finite_differences/differentiable.py

@@ -755,11 +758,12 @@ def __new__(cls, expr, mapper, **kwargs):
        obj = super().__new__(cls, expr, dimensions)
        obj._weights = weights
        obj._mapper = frozendict(mapper)
+        obj._order = order


comment on this?

I can actually probably drop it, it's now unused. It's the derivative order

mloubout · 2024-02-12T14:39:27Z

devito/ir/equations/algorithms.py


+def _lower_exprs(expressions, subs):


Note for later: subs should not be added to mapper as it would make compiler pass easier (pure sympy) and do the subs at the tail end of iet pass

mloubout · 2024-02-12T14:40:29Z

devito/passes/clusters/derivatives.py

-        e, clusters = _core(a, c, weights, reusables, mapper, sregistry)
-        args.append(e)
-        processed.extend(clusters)
+    if not isinstance(expr, IndexDerivative):


single dispatch on expr?

not gonna work unless we handle it here , but don't remember the details, will look again

ah no sorry, ignore the comment above, I had totally misunderstood it

Yes, I like the idea, will change tomorrow

done single-dispatch

EdCaunt · 2024-02-14T13:38:11Z

devito/core/cpu.py

@@ -63,9 +63,13 @@ def _normalize_kwargs(cls, **kwargs):
        # Distributed parallelism
        o['dist-drop-unwritten'] = oo.pop('dist-drop-unwritten', cls.DIST_DROP_UNWRITTEN)

-        # Misc
+        # Code generation options for derivatives


Could this be packed up into a function in a new file like devito.core.utils.py? It's exactly repeated in both cpu.py and gpu.py. Something like parse_cgen_options(o, oo)

probably yes, but you'll see there are a lot of small differences between the two

Not happening before Rice anyway, but I'll take note

FabioLuporini added the compiler label Feb 12, 2024

FabioLuporini requested a review from mloubout February 12, 2024 14:08

FabioLuporini commented Feb 12, 2024

View reviewed changes

mloubout reviewed Feb 12, 2024

View reviewed changes

EdCaunt reviewed Feb 14, 2024

View reviewed changes

FabioLuporini force-pushed the optimize-coeffs branch from 88c39ba to 7d76070 Compare February 20, 2024 09:32

FabioLuporini added 9 commits February 20, 2024 13:30

compiler: Tweak Weights and IndexDerivative

e986b95

compiler: Fix IndexDerivative lowering

27cd407

compiler: Decouple IndexDeriv into diff and nondiff

07f0990

compiler: Revamp and generalize lower_index_deriv

364bae0

compiler: Support derivatives unrolling

64343f1

compiler: Fix CDE

7f05a43

compiler: Use singledispatch to lower IndexDerivative

352cfa2

compiler: Remove useless IndexDerivative properties

261f38b

compiler: Patch dtype of lowered IndexDerivative

8db6b8b

FabioLuporini force-pushed the optimize-coeffs branch from 219b8c8 to 8db6b8b Compare February 20, 2024 13:31

mloubout approved these changes Feb 20, 2024

View reviewed changes

compiler: Apply user subs to IndexDerivatives

8f22f90

FabioLuporini merged commit ccfb823 into master Feb 22, 2024
31 checks passed

FabioLuporini deleted the optimize-coeffs branch February 22, 2024 13:44

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

compiler: Revamp lowering of IndexDerivatives #2310

compiler: Revamp lowering of IndexDerivatives #2310

FabioLuporini commented Feb 12, 2024

FabioLuporini Feb 12, 2024

codecov bot commented Feb 12, 2024 •

edited

Loading

mloubout Feb 12, 2024

FabioLuporini Feb 12, 2024 •

edited

Loading

mloubout Feb 12, 2024

mloubout Feb 12, 2024

FabioLuporini Feb 19, 2024

FabioLuporini Feb 20, 2024

mloubout Feb 12, 2024

FabioLuporini Feb 12, 2024

FabioLuporini Feb 20, 2024

FabioLuporini Feb 20, 2024

mloubout Feb 20, 2024

mloubout Feb 12, 2024

FabioLuporini Feb 12, 2024

FabioLuporini Feb 20, 2024

mloubout Feb 12, 2024

mloubout Feb 12, 2024

FabioLuporini Feb 12, 2024

FabioLuporini Feb 19, 2024

FabioLuporini Feb 20, 2024

EdCaunt Feb 14, 2024 •

edited

Loading

FabioLuporini Feb 19, 2024

		@@ -803,9 +815,13 @@ def _evaluate(self, **kwargs):
		return EvalDerivative(expr, base=self.base)


		class DiffDerivative(IndexDerivative, DifferentiableOp):

compiler: Revamp lowering of IndexDerivatives #2310

compiler: Revamp lowering of IndexDerivatives #2310

Conversation

FabioLuporini commented Feb 12, 2024

Choose a reason for hiding this comment

codecov bot commented Feb 12, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

FabioLuporini Feb 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

EdCaunt Feb 14, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Feb 12, 2024 •

edited

Loading

FabioLuporini Feb 12, 2024 •

edited

Loading

EdCaunt Feb 14, 2024 •

edited

Loading