Skip to content

Commit

Permalink
SCCRawStack: Re-organise import and make static helper a property
Browse files Browse the repository at this point in the history
  • Loading branch information
mlange05 authored and reuterbal committed May 10, 2024
1 parent 6a8640d commit 9f9ca78
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 28 deletions.
42 changes: 24 additions & 18 deletions loki/transformations/raw_stack_allocator.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,23 @@

import re

from loki.batch import Transformation
from loki.expression import Array, Scalar
from loki.types import BasicType
from loki.analyse import dataflow_analysis_attached
from loki.types import SymbolAttributes
from loki.ir import Assignment, CallStatement, Pragma
from loki.tools import as_tuple
from loki.ir import FindNodes, Transformer
from loki.batch import ProcedureItem
from loki.backend.fgen import fgen
from loki.batch.item import ProcedureItem
from loki.batch.transformation import Transformation
from loki.expression.symbols import (
Array, Scalar, Variable, Literal, Product, Sum, InlineCall,
IntLiteral, RangeIndex, DeferredTypeSymbol
)
from loki.expression.symbolic import is_dimension_constant, simplify
from loki.expression.mappers import DetachScopesMapper
from loki.expression.expr_visitors import FindVariables, SubstituteExpressions
from loki.backend.fgen import fgen
from loki.expression.symbols import (
Variable, Literal, Product, Sum, InlineCall, IntLiteral, RangeIndex, DeferredTypeSymbol)
from loki.ir.nodes import Assignment, CallStatement, Pragma
from loki.ir.find import FindNodes
from loki.ir.transformer import Transformer
from loki.tools import as_tuple
from loki.types import BasicType, SymbolAttributes


__all__ = ['TemporariesRawStackTransformation']

Expand Down Expand Up @@ -74,6 +76,12 @@ class TemporariesRawStackTransformation(Transformation):
# Traverse call tree in reverse when using Scheduler
reverse_traversal = True

type_name_dict = {
BasicType.REAL: {'kernel': 'P', 'driver': 'Z'},
BasicType.LOGICAL: {'kernel': 'LD', 'driver': 'LL'},
BasicType.INTEGER: {'kernel': 'K', 'driver': 'I'}
}

def __init__(self, block_dim, horizontal,
stack_name='STACK',
local_int_var_name_pattern='JD_{name}',
Expand All @@ -88,13 +96,11 @@ def __init__(self, block_dim, horizontal,
if key:
self._key = key


int_type = SymbolAttributes(dtype=BasicType.INTEGER, kind=DeferredTypeSymbol('JPIM'))

type_name_dict = {BasicType.REAL: {'kernel': 'P', 'driver': 'Z'},
BasicType.LOGICAL: {'kernel': 'LD', 'driver': 'LL'},
BasicType.INTEGER: {'kernel': 'K', 'driver': 'I'}}

@property
def int_type(self):
return SymbolAttributes(
dtype=BasicType.INTEGER, kind=DeferredTypeSymbol('JPIM')
)

def transform_subroutine(self, routine, **kwargs):

Expand Down
18 changes: 8 additions & 10 deletions loki/transformations/tests/test_raw_stack_allocator.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,20 +9,18 @@

import pytest

from loki.tools import gettempdir
from loki.dimension import Dimension
from loki.backend import fgen
from loki.batch import Scheduler, SchedulerConfig, ProcedureItem
from loki.frontend.util import OMNI
from loki.backend.fgen import fgen
from loki.types import BasicType
from loki.dimension import Dimension
from loki.expression import DeferredTypeSymbol, InlineCall, IntLiteral
from loki.frontend import available_frontends, OMNI
from loki.ir import FindNodes, CallStatement, Assignment, Pragma
from loki.sourcefile import Sourcefile
from loki.expression.symbols import DeferredTypeSymbol, InlineCall, IntLiteral
from loki.transform.transform_array_indexing import normalize_range_indexing

from conftest import available_frontends
from loki.tools import gettempdir
from loki.types import BasicType

from transformations.raw_stack_allocator import TemporariesRawStackTransformation
from loki.transformations.array_indexing import normalize_range_indexing
from loki.transformations.raw_stack_allocator import TemporariesRawStackTransformation


@pytest.fixture(scope='module', name='block_dim')
Expand Down

0 comments on commit 9f9ca78

Please sign in to comment.