Const Assignment Map Fusion: If two maps assign the same value for every element in a subset of the underlying array (and the subset is not dependent on the array in any way), then we can often fuse the two maps (not always possible) #1685

pratyai · 2024-10-15T12:29:42Z

Since scheduling multiple map kernels with very little internal operation can have a large overhead, sometimes we would like to fuse two such maps if possible. Constant assignment maps are such a case. If two maps assign the same value for every element in a subset of the underlying array (and the subset is not dependent on the array in any way), then:

If the two maps' subsets are identical, we can simply fuse the two maps by moving the body of one map to another (with appropriate wiring), since we know the order of the assignments do not matter here, and can even be deduplicated in some cases.
If the two maps' subsets are not identical, even then we can occasionally fuse them using a grid-strided loop pattern (which essentially emulates a conditional to ensure that only the appropriate elements are assigned).

Motivating Example

Consider the following graphs, all representing a computation that assigns 1 to the boundary of a 2D domain. The first table represents the graphs scheduled for CPU, the second for GPU.

Device	Original	w/o GSL	with GSL
CPU
GPU

Performance

We have profiled a 2D and a 3D boundary initialization, both on CPU and GPU (both on Davinci cluster).
Benchmark scripts and reports are to be found in https://github.com/pratyai/dace/tree/bench-const-assignment-fusion
I will be quoting the performance summaries in further comments.

Comment on GPU performance

The GPU transformation adds additional operation copying the entire array to and from GPU memory, resulting in O(n^d) main <=> GPU movement, whereas the assignment itself only touches O(n^{d-1}) elements. However, this is because the benchmark itself does not do anything else but the assignment. In real computations, we would likely need to move the entire array anyway.

Because of this, it is probably better to just focus on the combined performance of the map kernels here.

pratyai · 2024-10-15T12:30:41Z

CPU Results (comment might be updated)

benchmark_const_assignment_fusion_test_assign_top_row_0 0.13079101336188614 ms
===2D boundary init: original op===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              0.108          0.128          0.131          0.366          
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_top_row_47)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.007          0.020          0.021          0.045          
---------------------------------------------------------------------------
|-State (1)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_bottom_row_53)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.005          0.006          0.006          0.010          
---------------------------------------------------------------------------
|-State (2)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_left_col_59)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.004          0.005          0.005          0.024          
---------------------------------------------------------------------------
|-State (3)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_right_col_65)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.004          0.006          0.005          0.100          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_top_row_0 0.11151551734656096 ms
===2D boundary init: fused op w/o. grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              0.090          0.111          0.112          0.920          
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (4, Map map_fusion_wrapper)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.014          0.021          0.018          0.151          
---------------------------------------------------------------------------
| |-Node (9, Map map_fusion_wrapper)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.004          0.006          0.005          0.054          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_top_row_0 0.10572001338005066 ms
===2D boundary init: fused op with grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              0.084          0.103          0.106          0.393          
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (8, Map map_fusion_wrapper)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.009          0.022          0.022          0.024          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_bounary_3d 0.23336600861512125 ms
===3D boundary init: original op===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              0.214          0.233          0.233          0.497          
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_top_face_90)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.020          0.027          0.027          0.054          
---------------------------------------------------------------------------
|-State (1)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_bottom_face_9)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.013          0.015          0.014          0.038          
---------------------------------------------------------------------------
|-State (2)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_front_face_10)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.012          0.013          0.013          0.035          
---------------------------------------------------------------------------
|-State (3)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_back_face_108)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.010          0.011          0.011          0.038          
---------------------------------------------------------------------------
|-State (4)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_left_face_114)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.029          0.031          0.031          0.051          
---------------------------------------------------------------------------
|-State (5)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_right_face_12)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.024          0.026          0.025          0.051          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_bounary_3d 0.19121551304124296 ms
===3D boundary init: fused op w/o. grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              0.173          0.192          0.191          0.427          
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (4, Map map_fusion_wrapper)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.026          0.033          0.033          0.056          
---------------------------------------------------------------------------
| |-Node (9, Map map_fusion_wrapper)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.018          0.020          0.020          0.082          
---------------------------------------------------------------------------
| |-Node (14, Map map_fusion_wrapper)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.037          0.043          0.042          0.066          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_bounary_3d 0.18826551968231797 ms
===3D boundary init: fused op with grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              0.170          0.188          0.188          0.469          
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (14, Map map_fusion_wrapper)                                                            
| | |Thread 12176263715017013721:                                                            
| | |          0.091          0.099          0.096          0.200          
---------------------------------------------------------------------------

pratyai · 2024-10-15T12:31:28Z

GPU Results (comment might be updated)

benchmark_const_assignment_fusion_test_assign_top_row_0 167.06380699179135 ms
===2D boundary init: original op===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              166.200        167.907        167.064        172.466        
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_top_row_47)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.002          0.004          0.004          0.009          
---------------------------------------------------------------------------
|-State (1)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_bottom_row_53)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.004          0.005          0.004          0.018          
---------------------------------------------------------------------------
|-State (2)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_left_col_59)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.004          0.006          0.006          0.014          
---------------------------------------------------------------------------
|-State (3)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_right_col_65)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.005          0.006          0.006          0.014          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_top_row_0 170.15381151577458 ms
===2D boundary init: fused op w/o. grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              166.765        169.892        170.154        172.003        
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (6, Map map_fusion_wrapper)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.003          0.004          0.004          0.008          
---------------------------------------------------------------------------
| |-Node (11, Map map_fusion_wrapper)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.008          0.009          0.008          0.022          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_top_row_0 165.11968348640949 ms
===2D boundary init: fused op with grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              164.683        165.159        165.120        167.484        
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (10, Map map_fusion_wrapper)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.007          0.008          0.008          0.013          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_bounary_3d 50.67180350306444 ms
===3D boundary init: original op===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              49.457         50.629         50.672         51.666         
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_top_face_90)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.008          0.008          0.008          0.013          
---------------------------------------------------------------------------
|-State (1)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_bottom_face_9)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.008          0.009          0.010          0.023          
---------------------------------------------------------------------------
|-State (2)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_front_face_10)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.008          0.008          0.008          0.024          
---------------------------------------------------------------------------
|-State (3)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_back_face_108)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.006          0.008          0.008          0.027          
---------------------------------------------------------------------------
|-State (4)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_left_face_114)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.042          0.044          0.044          0.053          
---------------------------------------------------------------------------
|-State (5)                                                                
| |-Node (0, Map benchmark_const_assignment_fusion_test_assign_right_face_12)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.044          0.045          0.045          0.055          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_bounary_3d 50.77531602000818 ms
===3D boundary init: fused op w/o. grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              50.484         50.773         50.775         51.235         
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (6, Map map_fusion_wrapper)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.008          0.008          0.008          0.012          
---------------------------------------------------------------------------
| |-Node (11, Map map_fusion_wrapper)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.006          0.008          0.008          0.020          
---------------------------------------------------------------------------
| |-Node (16, Map map_fusion_wrapper)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.050          0.052          0.052          0.065          
---------------------------------------------------------------------------


benchmark_const_assignment_fusion_test_assign_bounary_3d 49.49491852312349 ms
===3D boundary init: fused op with grid-strided loop===
Instrumentation report
SDFG Hash: 
---------------------------------------------------------------------------
Element        Runtime (ms)   
               Min            Mean           Median         Max            
---------------------------------------------------------------------------
SDFG (0)                                                                   
|:                                                                         
|              49.321         49.533         49.495         50.064         
---------------------------------------------------------------------------
|-State (0)                                                                
| |-Node (16, Map map_fusion_wrapper)                                                            
| | |Thread 3118465057716905996:                                                            
| | |          0.055          0.057          0.057          0.063          
---------------------------------------------------------------------------

pratyai · 2024-10-28T17:44:19Z

@ThrudPrimrose

The PR contains all the working ideas I had about the constant assignment map fusion:

Fusing two maps if they all "blindly" assign constants (for multiple arrays, all the ranges of that array must always get the same constant, so that the order does not matter)
Can fuse maps if the constant assignments are hidden behind a branch that only depends on loop parameters (so that we can fuse grid-strided-loops too).
Fusing two states too, if after that the map fusion is possible (as described before).
Optionally, can fuse if the ranges are not exactly the same using grid-strided-loop.

If it is too many feature at once, I can remove the non-essential pieces for future PRs too. Please let me know if you'd see more tests, different implementation or organization etc.

match.

depends on the map parameters.

grid-stride loop as a guard.

… one.

stride 1)

not happening)

maps, create empty ones.

hand-crafted graphs.

cover more cases in the state fusion.

pratyai force-pushed the const-assignment-fusion branch from 1c96bf4 to e1bf9a6 Compare October 24, 2024 11:30

pratyai added the no-ci Do not run any CI or actions for this PR label Oct 24, 2024

pratyai force-pushed the const-assignment-fusion branch from e1bf9a6 to db94444 Compare October 26, 2024 09:30

pratyai requested a review from ThrudPrimrose October 28, 2024 17:36

pratyai added 21 commits October 29, 2024 16:30

First draft: only within-state fusion, and only if the ranges exactly

d0a25f2

match.

add interstate fusion in there

0c76246

Handle the free floating maps too, whenever possible.

1ef574b

Handle the fact that loop variable names can be different.

5121081

handle if the const assignments are behind an if branch that only

95ea47a

depends on the map parameters.

small changes and documentation

e9be112

Don't use cache in tests.

4bc5d14

Fix the "no dependency" pattern + a small refactor.

c8206b1

if the ranges are not exactly the same, still try to fuse by adding a

378673a

grid-stride loop as a guard.

Fix the bound for grid strided loop.

06cfcab

Set the schedule type to the inner GSLs to sequential.

39ee843

Allow multiple nested maps, by starting to look at only the outermost…

d2e6f9f

… one.

Rewrite the fusion pattern to allow fusing nested maps if possible.

58b291f

move out the @staticmethods out of the class definiton.

b2f950b

Make the "grid-strided loop" a configurable option.

6ef22b7

Fix the range fusion (subsets.union() only gives a bounding box with

868266e

stride 1)

Add checks to make sure the transformations are actually happening (or

8568c9f

not happening)

If there was no path established between the parent and the strided

113d2b3

maps, create empty ones.

Be slightly more explicit about when to use grid strided loop.

e60fe6d

Handle the AST API differences between python versions.

cb5398d

Add two designated "negative tests".

a4e6071

pratyai force-pushed the const-assignment-fusion branch from 0cd919f to f2771d5 Compare October 29, 2024 15:36

Privatize the helper functions.

4e8ca45

Cover even more negative test cases.

23f2689

pratyai force-pushed the const-assignment-fusion branch from f2771d5 to 23f2689 Compare October 29, 2024 16:20

pratyai added 6 commits October 30, 2024 13:07

Construct a subgraph view to do the replacement correctly.

1167380

Replace all the convenient uses of @dace.program and simply() with

4d6b4c6

hand-crafted graphs.

Fix the "taking a function's reference instead of calling it" problem.

fb01a23

Flipped the map incorrectly when constructing the graph.

78f0c2f

Replace add_memlet_path() calls.

44e81ce

Simplify and clarify the no-dependency-pattern check. Add more tests to

e525ce4

cover more cases in the state fusion.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Const Assignment Map Fusion: If two maps assign the same value for every element in a subset of the underlying array (and the subset is not dependent on the array in any way), then we can often fuse the two maps (not always possible) #1685

Const Assignment Map Fusion: If two maps assign the same value for every element in a subset of the underlying array (and the subset is not dependent on the array in any way), then we can often fuse the two maps (not always possible) #1685

pratyai commented Oct 15, 2024 •

edited

Loading

pratyai commented Oct 15, 2024

pratyai commented Oct 15, 2024

pratyai commented Oct 28, 2024

Const Assignment Map Fusion: If two maps assign the same value for every element in a subset of the underlying array (and the subset is not dependent on the array in any way), then we can often fuse the two maps (not always possible) #1685

Are you sure you want to change the base?

Const Assignment Map Fusion: If two maps assign the same value for every element in a subset of the underlying array (and the subset is not dependent on the array in any way), then we can often fuse the two maps (not always possible) #1685

Conversation

pratyai commented Oct 15, 2024 • edited Loading

Motivating Example

Performance

Comment on GPU performance

pratyai commented Oct 15, 2024

CPU Results (comment might be updated)

pratyai commented Oct 15, 2024

GPU Results (comment might be updated)

pratyai commented Oct 28, 2024

pratyai commented Oct 15, 2024 •

edited

Loading