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 SCC wrapper transformation #124

Merged
merged 3 commits into from
Aug 29, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 47 additions & 16 deletions scripts/loki_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,6 @@
from transformations.utility_routines import DrHookTransformation, RemoveCallsTransformation
from transformations.pool_allocator import TemporariesPoolAllocatorTransformation
from transformations.single_column_claw import ExtractSCATransformation, CLAWTransformation
from transformations.single_column_coalesced_wrapper import SingleColumnCoalescedTransformation
from transformations.single_column_coalesced import (
SCCBaseTransformation, SCCAnnotateTransformation, SCCHoistTransformation
)
Expand Down Expand Up @@ -458,7 +457,15 @@ def ecphys(mode, config, header, source, build, cpp, directive, frontend):
# Create and setup the scheduler for bulk-processing
paths = [Path(s).resolve() for s in source]
paths += [Path(h).resolve().parent for h in header]
scheduler = Scheduler(paths=paths, config=config, definitions=definitions, frontend=frontend, preprocess=cpp)
scheduler = Scheduler(
paths=paths, config=config, definitions=definitions,
frontend=frontend, preprocess=cpp
)

# Pick the model dimensions out of the config file
horizontal = scheduler.config.dimensions.get('horizontal', None)
vertical = scheduler.config.dimensions.get('vertical', None)
block_dim = scheduler.config.dimensions.get('block_dim', None)

# Backward insert argument shapes (for surface routines)
scheduler.process(transformation=ArgumentArrayShapeAnalysis())
Expand All @@ -467,39 +474,63 @@ def ecphys(mode, config, header, source, build, cpp, directive, frontend):

# Remove DR_HOOK and other utility calls first, so they don't interfere with SCC loop hoisting
if 'scc' in mode:
scheduler.process(transformation=RemoveCallsTransformation(
utility_call_trafo = RemoveCallsTransformation(
routines=['DR_HOOK', 'ABOR1', 'WRITE(NULOUT'], include_intrinsics=True, kernel_only=True
))
)
else:
scheduler.process(transformation=DrHookTransformation(mode=mode, remove=False))
utility_call_trafo = DrHookTransformation(mode=mode, remove=False)
scheduler.process(transformation=utility_call_trafo)

# Now we instantiate our transformation pipeline and apply the main changes
# Instantiate transformation pipeline and apply the main changes
transformation = None
if mode == 'idem':
transformation = IdemTransformation()

if mode == 'sca':
# Define the target dimension to strip from kernel and caller
horizontal = scheduler.config.dimensions['horizontal']
transformation = ExtractSCATransformation(horizontal=horizontal)
Copy link
Contributor

Choose a reason for hiding this comment

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

Do we not need to wrap this in a tuple? Same for the idem trafo above.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Good catch! I've added an as_tuple around the transformations in the application loop, which will have the same effect, but makes it a little safer.


if mode in ['scc', 'scc-hoist']:
horizontal = scheduler.config.dimensions['horizontal']
vertical = scheduler.config.dimensions['vertical']
block_dim = scheduler.config.dimensions['block_dim']
transformation = SingleColumnCoalescedTransformation(
horizontal=horizontal, vertical=vertical, block_dim=block_dim,
directive=directive, hoist_column_arrays='hoist' in mode
# Compose the main SCC transformation from core components based on config
transformation = (
SCCBaseTransformation(horizontal=horizontal, directive=directive),
)

# Devectorise the single-column kernels by removing the horizotnal vector loops
transformation += (
SCCDevectorTransformation(horizontal=horizontal, trim_vector_sections=False),
)

# Demote vector temporaries to scalars where possible
transformation += (SCCDemoteTransformation(horizontal=horizontal),)

# Either hoist the vector loops to the driver or re-vectorise the kernels
if 'hoist' in mode:
transformation += (
SCCHoistTransformation(horizontal=horizontal, vertical=vertical, block_dim=block_dim),
)
else:
transformation += (SCCRevectorTransformation(horizontal=horizontal),)

# Add offload annotations depending on the desired programming model directives
transformation += (
SCCAnnotateTransformation(
horizontal=horizontal, vertical=vertical, directive=directive,
block_dim=block_dim, hoist_column_arrays='hoist' in mode
),
)

if transformation:
scheduler.process(transformation=transformation)
# Apply the set of trnasformations over the full depency tree in turn
for transform in as_tuple(transformation):
scheduler.process(transformation=transform)
else:
raise RuntimeError('[Loki] Convert could not find specified Transformation!')

# Apply the dependency-injection transformation
dependency = DependencyTransformation(mode='module', module_suffix='_MOD',
suffix=f'_{mode.upper()}')
dependency = DependencyTransformation(
mode='module', module_suffix='_MOD', suffix=f'_{mode.upper()}'
)
scheduler.process(transformation=dependency)

# Write out all modified source files into the build directory
Expand Down
Loading