Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove add_mul_fusion rewrite #1243

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

ricardoV94
Copy link
Contributor

Some cleanup that emerged while working on #1242

The point of this rewrite was to fuse nested add and mul operations into a single add/ mul. This however is always (and more completely) done by canonicalization before the rewrite is ever triggered. As evidence of this, removing this rewrite would not cause any of the Fusion related tests to fail.

I also added a direct rewrite for the desired behavior, of giving precedence to fusing add/mul before composites are introduced.

@@ -896,16 +896,16 @@ def print_profile(cls, stream, prof, level=0):
print(blanc, " time_toposort", prof[7], file=stream)


fuse_seqopt = SequenceDB()
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was potentially buggy, as fuse_seqopt was always being imported by the other module, even though it's creation was conditional on a config flag.

@@ -896,16 +896,16 @@ def print_profile(cls, stream, prof, level=0):
print(blanc, " time_toposort", prof[7], file=stream)


fuse_seqopt = SequenceDB()
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, should we get rid of it if we have only a single rewrite now?

@ricardoV94 ricardoV94 changed the title Remove add mul fusion Remove add_mul_fusion rewrite Oct 7, 2022
@codecov
Copy link

codecov bot commented Oct 7, 2022

Codecov Report

Merging #1243 (14697eb) into main (a619a4e) will increase coverage by 0.03%.
The diff coverage is 100.00%.

❗ Current head 14697eb differs from pull request most recent head 9e638e3. Consider uploading reports for the commit 9e638e3 to get more accurate results

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1243      +/-   ##
==========================================
+ Coverage   79.14%   79.17%   +0.03%     
==========================================
  Files         173      173              
  Lines       48528    48503      -25     
  Branches    10322    10315       -7     
==========================================
- Hits        38408    38403       -5     
+ Misses       7628     7613      -15     
+ Partials     2492     2487       -5     
Impacted Files Coverage Δ
aesara/tensor/rewriting/math.py 87.40% <ø> (+0.11%) ⬆️
aesara/tensor/rewriting/elemwise.py 85.75% <100.00%> (ø)
aesara/graph/basic.py 89.29% <0.00%> (+0.29%) ⬆️
aesara/sparse/basic.py 83.27% <0.00%> (+0.39%) ⬆️
aesara/sparse/rewriting.py 75.91% <0.00%> (+0.40%) ⬆️
aesara/typed_list/basic.py 89.79% <0.00%> (+1.36%) ⬆️

@ricardoV94 ricardoV94 marked this pull request as ready for review October 7, 2022 14:30
@ricardoV94 ricardoV94 requested review from brandonwillard and removed request for brandonwillard October 7, 2022 14:30
@ricardoV94 ricardoV94 marked this pull request as draft October 7, 2022 16:03
@ricardoV94
Copy link
Contributor Author

ricardoV94 commented Oct 7, 2022

Sometimes the local_add_canonizer and local_mul_canonizer fail to merge all nested operations. The pattern I notice is when we combine them with the inverse of the op:

import aesara
import aesara.tensor as at
from aesara.compile.mode import Mode

x, y, z = at.vectors("x", "y", "z")

mul = x * y * z
add = x + y + z

out = mul - add
f = aesara.function([x, y, z], [out], mode=Mode().including("canonicalize").excluding("fusion"))
aesara.dprint(f)
Elemwise{Sub}[(0, 0)] [id A] 3
 |Elemwise{mul,no_inplace} [id B] 2
 | |x [id C]
 | |y [id D]
 | |z [id E]
 |Elemwise{Add}[(0, 0)] [id F] 1
   |Elemwise{add,no_inplace} [id G] 0
   | |x [id C]
   | |y [id D]
   |z [id E]

Or, conversely

out = add / mul
f = aesara.function([x, y, z], [out], mode=Mode().including("canonicalize").excluding("fusion"))
aesara.dprint(f)
Elemwise{TrueDiv}[(0, 0)] [id A] 3
 |Elemwise{add,no_inplace} [id B] 2
 | |x [id C]
 | |y [id D]
 | |z [id E]
 |Elemwise{Mul}[(0, 0)] [id F] 1
   |Elemwise{mul,no_inplace} [id G] 0
   | |x [id C]
   | |y [id D]
   |z [id E]
Out[5]: <_io.TextIOWrapper name='<stdout>' mode='w' encoding='utf-8'>

@brandonwillard brandonwillard added graph rewriting refactor This issue involves refactoring tensor algebra Relates to our use and representations of tensor algebra labels Oct 7, 2022
@brandonwillard
Copy link
Member

brandonwillard commented Oct 7, 2022

Sometimes the local_add_canonizer and local_mul_canonizer fail to merge all nested operations. The pattern I notice is when we combine them with the inverse of the op:

Ha, I just commented on this in #1244; looks like I should've looked here first to better understand the context. Nevermind, it seems like this issue is fairly distinct from the main thrust of this PR.

@@ -896,16 +896,16 @@ def print_profile(cls, stream, prof, level=0):
print(blanc, " time_toposort", prof[7], file=stream)


fuse_seqopt = SequenceDB()
if config.tensor__local_elemwise_fusion:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't thought enough about this just yet, but it really looks like we need to consider removing this option entirely, because, unless we redesign the whole approach to fusion/FusionOptimizer, I don't know whether this idea of applying a graph (i.e. "global") rewriter on each node (i.e. "locally") makes any sense.

I could see a place for a distinct node-based fusion rewrite, but I really don't think it would/should operate anything like a graph-based rewrite (e.g. it seems like there would be way too much unnecessary graph walking).

In general, the idea of a node-based fusion rewrite sounds appealing, because it could be a lot simpler than the/a graph-based one, but I don't know which types of graphs the former would miss that the latter wouldn't.

The same job is done by canonicalize before this rewrite is ever called.
@brandonwillard brandonwillard changed the title Remove add_mul_fusion rewrite Remove add_mul_fusion rewrite Oct 11, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
graph rewriting refactor This issue involves refactoring tensor algebra Relates to our use and representations of tensor algebra
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants