feat[next]: Optimisations for icon4py

[samkellerhals]

Description

Changes to speedup diffusion granule execution from ICON.

• Removing isinstance checks from extract_connectivity_args and convert_args.
• Removing warning from field _maker.
• Allows connectivities to be passed (effectively allowing caching of connectivities on the caller side)

Further (future) optimisations:

• All connectivities are cached on the icon4py side, so the use of ensure_is_on_device is only necessary now in order to be able to run certain gt4py tests. This should be removed in the future (and done in the gt4py tests itself) as it adds an overhead here that is not necessary.
• convert_args still takes up the bulk of time, can we get rid of it or improve it?
• We have to pass array sizes for each dimension to gt4py stencils if no explicit domain bounds are defined, can this be standardised?

[havogt]

Forwarding the review to @egparedes for improving Python patterns in performance critical code parts.

[havogt]

Why this change? Can you undo?

[havogt]

I see it was intentional. Still should be undone and fixed on the icon4py side.

[samkellerhals]

So basically emitting the warning was slowing things down on the icon4py side when creating the gt4py fields when calling the granule. In what way could this be fixed on the icon4py side?

[samkellerhals]

The reason we currently still use np_as_located_field is because it gives you back a numpy array view of the pointer we pass from fortran. If we switch to as_field it does not use the same memory location any longer (I guess it makes a copy?) which should be fixed on the gt4py side.

[samkellerhals]

I propose to just wrap the warning in an if __debug__ for now, so that it can be disabled in -O mode.

[havogt]

Ok, I see, then it probably makes sense to use (temporarily) the common._field() function to wrap an array into a field.

[havogt]

feels like code should be refactored such that we always pass conn_args here. Probably makes sense to do this change in the context of a coarser refactoring. @DropD might have ideas.

[DropD]

Once I get size args extraction moved to this stage, I will look into adding that.

[havogt]

would functools.singledispatch work?

[samkellerhals]

I think it could work, however whilst nicer I have the feeling that it would be slower than a simple dictionary lookup.

[DropD]

I think it would be preferable to sacrifice readability only on the basis of hard evidence.

[havogt]

But in general, could make sense to discuss this in the context of going towards frozen programs.

[DropD]

If this pattern proves to be a significant optimization over singledispatch, it still needs to be made more readable.

I propose to encode the pattern in a class. That class then needs a docstring explaining when to use it over standard approaches and why, for example, subclasses of the relevant types won't work with it unless they are added explicitly to the dict.

Sketch:

class FastDispatch:
    """
    Optimized version of functools.singledispatch, does not take into account inheritance or protocol membership.
    
    This leads to a speed-up of XXX, as documented in ADR YYY.
    
    Usage:
    >>> @Fastdispatch.fastdispatch(Any)
    ... def extract_connectivity_args(connectivity, *args, **kwargs):
    ...     return None
    ...
    ... @extract_connectivity_args(NeighborTableOffsetProvider):
    ... def extract_connectivity_args_from_nbtable(connectivity, device, *args, **kwargs):
    ...     return (_ensure_is_on_device(connectivity.table, device), zero_tuple)
    """
    _registry: dict[type: str]
    
    def __call__(self, dispatch_arg, *args, **kwargs):
        return getattr(self, self._registry[type(dispatch_arg)])(dispatch_arg, *args, **kwargs)
        
    def register(self, type):
        def decorator(function):
            self._registry[type] = function
            return function
        return decorator
        
    @classmethod
    def fastdispatch(cls, default_type):
        return decorator(function):
            dispatcher = cls()
            dispatcher.register(default_type)(function)
            return dispatcher
        return decorator