Implement loop fusion optimization inside kernels #253

NaderAlAwar · 2024-01-25T15:59:05Z

This PR adds support for doing loop fusion inside a kernel. It can be enabled by setting the PK_LOOP_FUSE env variable. The conditions for two loops to be fused are as follows:

They must belong to the same scope OR their parent nodes are two for loops that can be fused.
They must have identical loop ranges. Right now this is checked by comparing the arguments of the two range calls as strings.
They must be adjacent OR the nodes between them are safe to move to before the first loop.
There must be no invalid dependencies (not implemented yet).

…looking at loop ranges, fix issue with looking at nodes between loops when the parent is the same, and add initial logic for dependency analysis

… C++)

…r memory operations fusion

test_loop_fusion.py: Cleaned up format

gliga · 2024-01-31T17:05:58Z

tests/test_loop_fusion.py

+        for k in range(2):
+            pk.printf("print 2 %d\n", k)
+
+def main():


can we format these as proper pytests so we can actually run them and have assertions in the end of each workunit?

gliga

Looks good except that comment about tests

gliga

Provided only comments on tests

loop_fusion_kernels.py: Contains bodies of kernels/tests test_loop_fusion.py: Test class to set environment, run, capture, and compare outputs of kernels

…cilitate fusion of memory ops

…nd on prior operation

…object

…ame too long errors

…ge kernels

NaderAlAwar added 5 commits January 25, 2024 09:45

Fusion: remove redundant str() call

76d7fbe

Fusion: move visitors to a separate file to reuse them in loop fusion

6af0b2a

Translators: add index in body to each AST node

8fd27c3

Optimizations: add the loop fusion optimization

a1b3633

Compiler: add env variable to do loop fusion

778dff7

NaderAlAwar requested a review from gliga January 25, 2024 15:59

NaderAlAwar marked this pull request as draft January 25, 2024 15:59

NaderAlAwar and others added 9 commits January 27, 2024 09:11

Translators: add idx in parent body attribute to all nodes

c266301

Optimizations: account for fused_* additions from fused kernels when …

74f5cdd

…looking at loop ranges, fix issue with looking at nodes between loops when the parent is the same, and add initial logic for dependency analysis

PyKokkosVisitors: add pk.cpp_auto as an allowed type (maps to auto in…

dae698b

… C++)

StaticTranslator: fix idx_in_loop to work for all AST nodes with lists

f8060eb

Optimizations: move common functionality to util.py in preparation fo…

61ab4c7

…r memory operations fusion

Optimizations: add memory_ops_fuse optimization

c10f7f4

Compiler: restrict optimizations to workunit and fused styles only

a8928a8

test_loop_fusion.py: Added basic and manual tests for loop fusion

9eb72b4

run_loop_fusion_test.py: added basic automation and comparison

3563f68

test_loop_fusion.py: Cleaned up format

gliga reviewed Jan 31, 2024

View reviewed changes

gliga approved these changes Jan 31, 2024

View reviewed changes

gliga reviewed Jan 31, 2024

View reviewed changes

HannanNaeem and others added 5 commits January 31, 2024 12:58

cleaned up and refactored tests to work with pytest

e7295e0

loop_fusion_kernels.py: Contains bodies of kernels/tests test_loop_fusion.py: Test class to set environment, run, capture, and compare outputs of kernels

Added simple neg dist test

5ec2bc0

Optimizations: track constant variables used in memory accesses to fa…

ee61112

…cilitate fusion of memory ops

Fusion: implent rsub() for future and add check if data does not depe…

6dad848

…nd on prior operation

mypy: add missing ignore

c103ea5

NaderAlAwar marked this pull request as ready for review February 13, 2024 00:10

NaderAlAwar added 5 commits February 12, 2024 19:38

mypy: add missing ignore

53a2f44

StaticTranslator: fix issue with assigning parent_accessor to string …

edc53f0

…object

Fusion: overload more operators for Future

89186d5

Fusion: truncate long kernel names to avoid errors

3c733be

ModuleSetup: truncate long file names in fused kernels to avoid filen…

faa9bc0

…ame too long errors

NaderAlAwar added 3 commits February 16, 2024 12:13

WorkunitVisitor: fix issue with getting accumulator argument in MDRan…

95ccc6e

…ge kernels

Tests: fix formatting and spelling

76cc291

Optimizations: don't fuse loops with print statements

65cbeeb

NaderAlAwar merged commit 7bfd51b into kokkos:main Feb 19, 2024
5 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement loop fusion optimization inside kernels #253

Implement loop fusion optimization inside kernels #253

NaderAlAwar commented Jan 25, 2024

gliga Jan 31, 2024

gliga left a comment

gliga left a comment

Implement loop fusion optimization inside kernels #253

Implement loop fusion optimization inside kernels #253

Conversation

NaderAlAwar commented Jan 25, 2024

gliga Jan 31, 2024

Choose a reason for hiding this comment

gliga left a comment

Choose a reason for hiding this comment

gliga left a comment

Choose a reason for hiding this comment