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

Improve schedule/as_matmul matching on hybrid loop-libop code #522

Open
roastduck opened this issue Jul 7, 2023 · 1 comment
Open

Improve schedule/as_matmul matching on hybrid loop-libop code #522

roastduck opened this issue Jul 7, 2023 · 1 comment
Labels
enhancement New feature or request

Comments

@roastduck
Copy link
Owner

Current schedule/as_matmul can't match some basic cases when there are both loops and libop calls. Example:

for i in range(500):
  c[i] = a[i] @ ft.transpose(b, (1, 0))

This code cannot be mapped to a Matmul node because libop introduces intermediate local variables. The code is actually like this:

for i in range(500):
  t = ft.transpose(b, (1, 0))
  u = a[i] @ t
  c[i] = u

In order to deal with this case, we need two following changes:

  • Call inline on every variables inside the matching sub-tree. There will be no side-effect because the matching will fail otherwise. This will solve the problem of t in this example.
  • Implement a new schedule, maybe named anti_inline, that removes an intermediate variable an redirect all Stores to it to the final destination. We can implement anti_inline for only variables that are copied to another variable with modifications (like c[i] = u). These will also be no side-effect because the matching will fail otherwise. This will solve the problem of u.

After implementing the changes above, schedule/as_matmul is expected to match this code to a Matmul, but still can't deal with its derivatives. In order to accept such a code in AD, will need one more change:

  • Simultaneously match multiple Matmuls in one schedule/as_matmul call, instead of relying on #! prefer_libs to fission the loops.
@roastduck roastduck added the enhancement New feature or request label Jul 7, 2023
@roastduck
Copy link
Owner Author

Another approach to solve the problem (without considering AD) is to add a new schedule uncache and run it automatically inside as_matmul. As the name suggests, uncache undoes the cache schedule. uncache(v) detects whether v maps to a parent VarDef u, and replace v by u with corresponding indices.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

1 participant