SCCRawStackTransformation: Small fixes to the transformation itself

ecmwf-ifs · Jun 14, 2024 · d423585 · d423585
1 parent 9d3732f
commit d423585
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/loki/transformations/raw_stack_allocator.py b/loki/transformations/raw_stack_allocator.py
@@ -88,7 +88,7 @@ class TemporariesRawStackTransformation(Transformation):
     def __init__(
             self, block_dim, horizontal, stack_name='STACK',
             local_int_var_name_pattern='JD_{name}', directive=None,
-            key=None, driver_horizontal=None, **kwargs
+            driver_horizontal=None, **kwargs
     ):
         super().__init__(**kwargs)
         self.block_dim = block_dim
@@ -98,9 +98,6 @@ def __init__(
         self.directive = directive
         self.driver_horizontal = driver_horizontal
 
-        if key:
-            self._key = key
-
     @property
     def int_type(self):
         return SymbolAttributes(
@@ -151,7 +148,7 @@ def insert_stack_in_calls(self, routine, stack_arg_dict, successors):
 
         Parameters
         ----------
-        routine : :any:'Subroutine
+        routine : :any:`Subroutine`
             The routine in which to transform call statements
         stack_arg_dict : dict
             dict that maps dtype and kind to the sets of stack size variables
@@ -160,7 +157,7 @@ def insert_stack_in_calls(self, routine, stack_arg_dict, successors):
             The items corresponding to successor routines called from :data:`routine`
         """
         successor_map = {
-            successor.ir.name.lower(): successor
+            successor.local_name: successor
             for successor in successors if isinstance(successor, ProcedureItem)
         }
         call_map = {}
@@ -205,7 +202,7 @@ def create_stacks_driver(self, routine, stack_dict, successors):
 
         Parameters
         ----------
-        routine : :any:'Subroutine
+        routine : :any:`Subroutine`
             The driver subroutine to get the stack_variables
         stack_dict : dict
             dict that maps dtype and kind to an expression for the required stack size
@@ -385,7 +382,7 @@ def apply_raw_stack_allocator_to_temporaries(self, routine, item=None):
 
         Parameters
         ----------
-        routine : :any:'Subroutine
+        routine : :any:`Subroutine`
             Subroutine object to apply transformation to
 
         Returns