From e57ee723d85233eb81c3fc5af1efe2d73b40aab3 Mon Sep 17 00:00:00 2001
From: Kaushik Kulkarni <kaushikcfd@gmail.com>
Date: Sun, 18 Mar 2018 01:10:53 -0500
Subject: [PATCH] arg_id_to_descr is working
---
loopy/kernel/__init__.py | 6 +-
loopy/kernel/function_interface.py | 174 +++++++++++++++++++++++++----
loopy/library/function.py | 5 -
loopy/preprocess.py | 168 ++++++++++++++++++++++++----
loopy/symbolic.py | 13 ++-
loopy/type_inference.py | 100 +----------------
6 files changed, 316 insertions(+), 150 deletions(-)
diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py
index 851626a8d..d716f0b78 100644
--- a/loopy/kernel/__init__.py
+++ b/loopy/kernel/__init__.py
@@ -37,8 +37,7 @@ from pytools import UniqueNameGenerator, generate_unique_names
from loopy.library.function import (
default_function_mangler,
- single_arg_function_mangler,
- default_function_identifiers)
+ single_arg_function_mangler)
from loopy.diagnostic import CannotBranchDomainTree, LoopyError
from loopy.tools import natsorted
@@ -271,8 +270,7 @@ class LoopKernel(ImmutableRecordWithoutPickling):
# Populating the function identifiers based on the target and the default
# function identifiers
- function_identifiers = (default_function_identifiers() |
- target.get_device_ast_builder().function_identifiers())
+ function_identifiers = target.get_device_ast_builder().function_identifiers()
ImmutableRecordWithoutPickling.__init__(self,
domains=domains,
diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py
index 4bc7f3d76..7127d142b 100644
--- a/loopy/kernel/function_interface.py
+++ b/loopy/kernel/function_interface.py
@@ -1,11 +1,18 @@
from __future__ import division, absolute_import
+import re
+import six
import numpy as np
from pytools import ImmutableRecord
from loopy.diagnostic import LoopyError
from loopy.types import NumpyType
+from loopy.kernel.instruction import (MultiAssignmentBase, CInstruction,
+ _DataObliviousInstruction)
+
+from loopy.symbolic import IdentityMapper, ScopedFunction
+
# {{{ argument descriptors
@@ -21,17 +28,20 @@ class ArgDescriptor(ImmutableRecord):
mem_scope=None,
shape=None,
dim_tags=None):
- super(ArgDescriptor).__init__(self,
- mem_scope=mem_scope,
+ super(ArgDescriptor, self).__init__(mem_scope=mem_scope,
shape=shape,
dim_tags=dim_tags)
class ValueArgDescriptor(ArgDescriptor):
- """
- """
def __init__(self):
- super(ValueArgDescriptor, self).__init__(self)
+ super(ValueArgDescriptor, self).__init__()
+
+ def __str__(self):
+ return "ValueArgDescriptor"
+
+ def __repr__(self):
+ return "ValueArgDescriptor"
class ArrayArgDescriptor(ArgDescriptor):
@@ -41,9 +51,10 @@ class ArrayArgDescriptor(ArgDescriptor):
"""
def __init__(self,
+ shape=None,
mem_scope=None,
dim_tags=None):
- super(ArgDescriptor, self).__init__(self,
+ super(ArgDescriptor, self).__init__(shape=None,
mem_scope=mem_scope,
dim_tags=dim_tags)
@@ -266,10 +277,7 @@ class InKernelCallable(ImmutableRecord):
# {{{ attempt to specialize using scalar functions
- from loopy.library.function import default_function_identifiers
- if self.name in default_function_identifiers():
- ...
- elif self.name in target.get_device_ast_builder().function_identifiers():
+ if self.name in target.get_device_ast_builder().function_identifiers():
from loopy.target.c import CTarget
from loopy.target.opencl import OpenCLTarget
from loopy.target.pyopencl import PyOpenCLTarget
@@ -371,7 +379,36 @@ class InKernelCallable(ImmutableRecord):
its keyword identifier.
"""
- raise NotImplementedError()
+ if self.subkernel is None:
+ # This is a scalar call
+ # need to assert that the name is in funtion indentifiers
+ arg_id_to_descr[-1] = ValueArgDescriptor()
+ return self.copy(arg_id_to_descr=arg_id_to_descr)
+
+ else:
+ # Now this ia a kernel call
+ # tuning the subkernel so that we have the the matching shapes and
+ # dim_tags.
+ # FIXME: Although We receive input if the argument is
+ # local/global. We do not use it to set the subkernel function
+ # signature. Need to do it, so that we can handle teporary inputs
+ # in the array call.
+
+ # Collecting the parameters
+ new_args = self.args.copy()
+ kw_to_pos, pos_to_kw = get_kw_pos_association(self.subkernel)
+
+ for id, descr in arg_id_to_descr.items():
+ if isinstance(id, str):
+ id = kw_to_pos[id]
+ assert isinstance(id, int)
+ new_args[id] = new_args[id].copy(shape=descr.shape,
+ dim_tags=descr.dim_tags)
+
+ descriptor_specialized_knl = self.subkernel.copy(args=new_args)
+
+ return self.copy(subkernel=descriptor_specialized_knl,
+ arg_id_to_descr=arg_id_to_descr)
def with_iname_tag_usage(self, unusable, concurrent_shape):
"""
@@ -390,16 +427,10 @@ class InKernelCallable(ImmutableRecord):
raise NotImplementedError()
- def is_arg_written(self, arg_id):
- """
- :arg arg_id: (keyword) name or position
- """
-
- raise NotImplementedError()
-
def is_ready_for_code_gen(self):
- raise NotImplementedError()
+ return (self.arg_id_to_dtype is not None and
+ self.arg_id_to_descr is not None)
# {{{ code generation
@@ -413,6 +444,8 @@ class InKernelCallable(ImmutableRecord):
raise NotImplementedError()
def emit_call(self, target):
+ # two varieties of this call, when obtained in between a function and
+ # when obtained as a separate instruction statement.
raise NotImplementedError()
@@ -421,7 +454,7 @@ class InKernelCallable(ImmutableRecord):
def __eq__(self, other):
return (self.name == other.name
and self.arg_id_to_descr == other.arg_id_to_descr
- and self.arg_id_to_dtype == other.arg_id_to_keyword)
+ and self.arg_id_to_dtype == other.arg_id_to_dtype)
def __hash__(self):
return hash((self.name, self.subkernel))
@@ -530,4 +563,105 @@ class CallableKernel(InKernelCallable):
# }}}
+
+# {{{ new pymbolic calls to scoped functions
+
+def next_indexed_name(name):
+ func_name = re.compile(r"^(?P<alpha>\S+?)_(?P<num>\d+?)$")
+
+ match = func_name.match(name)
+
+ if match is None:
+ if name[-1] == '_':
+ return "{old_name}0".format(old_name=name)
+ else:
+ return "{old_name}_0".format(old_name=name)
+
+ return "{alpha}_{num}".format(alpha=match.group('alpha'),
+ num=int(match.group('num'))+1)
+
+
+class FunctionScopeChanger(IdentityMapper):
+ #TODO: Make it sophisticated as in I don't like the if-else systems. Needs
+ # something else.
+ def __init__(self, new_names):
+ self.new_names = new_names
+ self.new_names_set = frozenset(new_names.values())
+
+ def map_call(self, expr):
+ if expr in self.new_names:
+ return type(expr)(
+ ScopedFunction(self.new_names[expr]),
+ tuple(self.rec(child)
+ for child in expr.parameters))
+ else:
+ return IdentityMapper.map_call(self, expr)
+
+ def map_call_with_kwargs(self, expr):
+ if expr in self.new_names:
+ return type(expr)(
+ ScopedFunction(self.new_names[expr]),
+ tuple(self.rec(child)
+ for child in expr.parameters),
+ dict(
+ (key, self.rec(val))
+ for key, val in six.iteritems(expr.kw_parameters))
+ )
+ else:
+ return IdentityMapper.map_call_with_kwargs(self, expr)
+
+
+def register_pymbolic_calls_to_knl_callables(kernel,
+ pymbolic_calls_to_knl_callables):
+ """ Takes in a mapping :arg:`pymbolic_calls_to_knl_callables` and returns a
+ new kernel which includes an association with the given pymbolic calls to
+ instances of :class:`InKernelCallable`
+ """
+
+ scoped_names_to_functions = kernel.scoped_functions.copy()
+
+ # A dict containing the new scoped functions to the names which have been
+ # assigned to them
+ scoped_functions_to_names = {}
+
+ # A dict containing the new name that need to be assigned to the
+ # corresponding pymbolic call
+ pymbolic_calls_to_new_names = {}
+
+ for pymbolic_call, in_knl_callable in pymbolic_calls_to_knl_callables.items():
+ # checking if such a in-kernel callable already exists.
+ if in_knl_callable not in scoped_functions_to_names:
+ # No matching in_knl_callable found => make a new one with a new
+ # name.
+
+ unique_name = next_indexed_name(pymbolic_call.function.name)
+ while unique_name in scoped_names_to_functions:
+ # keep on finding new names till one a unique one is found.
+ unique_name = next_indexed_name(unique_name)
+
+ # book-keeping of the functions and names mappings for later use
+ scoped_names_to_functions[unique_name] = in_knl_callable
+ scoped_functions_to_names[in_knl_callable] = unique_name
+
+ pymbolic_calls_to_new_names[pymbolic_call] = (
+ scoped_functions_to_names[in_knl_callable])
+
+ # Using the data populated in pymbolic_calls_to_new_names to change the
+ # names of the scoped functions of all the calls in the kernel.
+ new_insns = []
+ scope_changer = FunctionScopeChanger(pymbolic_calls_to_new_names)
+ for insn in kernel.instructions:
+ if isinstance(insn, (MultiAssignmentBase, CInstruction)):
+ expr = scope_changer(insn.expression)
+ new_insns.append(insn.copy(expression=expr))
+ elif isinstance(insn, _DataObliviousInstruction):
+ new_insns.append(insn)
+ else:
+ raise NotImplementedError("Type Inference Specialization not"
+ "implemented for %s instruciton" % type(insn))
+ return kernel.copy(scoped_functions=scoped_names_to_functions,
+ instructions=new_insns)
+
+# }}}
+
# vim: foldmethod=marker
diff --git a/loopy/library/function.py b/loopy/library/function.py
index e8e1e22fa..3573f1d54 100644
--- a/loopy/library/function.py
+++ b/loopy/library/function.py
@@ -23,11 +23,6 @@ THE SOFTWARE.
"""
-def default_function_identifiers():
- from loopy.library.reduction import reduction_function_identifiers
- return set("make_tuple") | reduction_function_identifiers()
-
-
def default_function_mangler(kernel, name, arg_dtypes):
from loopy.library.reduction import reduction_function_mangler
diff --git a/loopy/preprocess.py b/loopy/preprocess.py
index d7d961d25..741f828e2 100644
--- a/loopy/preprocess.py
+++ b/loopy/preprocess.py
@@ -27,6 +27,7 @@ import six
from loopy.diagnostic import (
LoopyError, WriteRaceConditionWarning, warn_with_kernel,
LoopyAdvisory)
+from functools import reduce
import islpy as isl
@@ -37,11 +38,11 @@ from loopy.version import DATA_MODEL_VERSION
from loopy.kernel.data import make_assignment
# for the benefit of loopy.statistics, for now
from loopy.type_inference import infer_unknown_types
-from loopy.symbolic import ScopedFunction, IdentityMapper
+from loopy.symbolic import ScopedFunction, CombineMapper
from pymbolic.mapper import Collector
from loopy.kernel.instruction import (MultiAssignmentBase, CInstruction,
- _DataObliviousInstruction)
+ CallInstruction, _DataObliviousInstruction)
import logging
logger = logging.getLogger(__name__)
@@ -2127,38 +2128,155 @@ def check_functions_are_scoped(kernel):
# {{{ arg_descr_inference
-# take help from the work we did yesterday to populate this
-class ArgDescriptionAdder(IdentityMapper):
+def get_arg_description_from_sub_array_ref(sub_array, kernel):
+ """ Gets the dim_tags, memory scope, shape informations of a
+ :class:`SubArrayRef` argument in the caller kernel packed into
+ :class:`ArrayArgDescriptor`.
+ """
+ from loopy.kernel.function_interface import ArrayArgDescriptor
- def __init__(self,):
- ...
+ name = sub_array.subscript.attribute.name
- def map_call(self, expr):
- ...
+ if name in kernel.temporary_variables:
+ mem_scope = "LOCAL"
+ arg = kernel.temporary_variables[name]
+ assert name not in kernel.arg_dict
+ else:
+ assert name in kernel.arg_dict
+ mem_scope = "GLOBAL"
+ arg = kernel.arg_dict[name]
+
+ sub_dim_tags, sub_shape = sub_array.get_sub_array_dim_tags_and_shape(
+ arg.dim_tags, arg.shape)
+ return ArrayArgDescriptor(mem_scope=mem_scope,
+ dim_tags=sub_dim_tags,
+ shape=sub_shape)
-def arg_descr_inference(kernel):
+
+class ArgDescriptionInferer(CombineMapper):
+ """ Returns a set with elements as instances of :class:`tuple` (expr,
+ in_kenrel_callable). The mapped `in_kenrel_callable` of the
+ :class:`InKernelCallable` are descriptor specialized for the given
+ arguments.
+ """
+
+ def __init__(self, scoped_functions):
+ self.scoped_functions = scoped_functions
+
+ def combine(self, values):
+ import operator
+ return reduce(operator.or_, values, set())
+
+ def map_call(self, expr, **kwargs):
+ from loopy.kernel.function_interface import ValueArgDescriptor
+ from loopy.symbolic import SubArrayRef
+
+ # descriptors for the args
+ arg_id_to_descr = dict((i, get_arg_description_from_sub_array_ref(par))
+ if isinstance(par, SubArrayRef) else (i, ValueArgDescriptor())
+ for i, par in enumerate(expr.parameters))
+
+ assignee_id_to_descr = {}
+
+ # assignee descriptor
+ if 'assignees' in kwargs:
+ # If supplied with assignees then this is a CallInstruction
+ assignees = kwargs['assignees']
+ assert isinstance(assignees, tuple)
+ for i, par in enumerate(assignees):
+ if isinstance(par, SubArrayRef):
+ assignee_id_to_descr[-i-1] = (
+ get_arg_description_from_sub_array_ref(par))
+ else:
+ assignee_id_to_descr[-i-1] = ValueArgDescriptor()
+
+ # gathering all the descriptors
+ combined_arg_id_to_dtype = {**arg_id_to_descr, **assignee_id_to_descr}
+
+ # specializing the function according to the parameter description
+ new_scoped_function = (
+ self.scoped_functions[expr.function.name].with_descrs(
+ combined_arg_id_to_dtype))
+
+ # collecting the descriptors for args, kwargs, assignees
+ return set(((expr, new_scoped_function),))
+
+ def map_call_with_kwargs(self, expr, **kwargs):
+ from loopy.kernel.function_intergace import ValueArgDescriptor
+ from loopy.symbolic import SubArrayRef
+
+ # descriptors for the args and kwargs:
+ arg_id_to_descr = dict((i, get_arg_description_from_sub_array_ref(par))
+ if isinstance(par, SubArrayRef) else ValueArgDescriptor()
+ for i, par in enumerate(expr.parameters) +
+ expr.kw_parameters.items())
+
+ assignee_id_to_descr = {}
+
+ if 'assignees' in kwargs:
+ # If supplied with assignees then this is a CallInstruction
+ assignees = kwargs['assignees']
+ assert isinstance(assignees, tuple)
+ for i, par in enumerate(assignees):
+ if isinstance(par, SubArrayRef):
+ assignee_id_to_descr[-i-1] = (
+ get_arg_description_from_sub_array_ref(par))
+ else:
+ assignee_id_to_descr[-i-1] = ValueArgDescriptor()
+
+ # gathering all the descriptors
+ combined_arg_id_to_descr = {**arg_id_to_descr, **assignee_id_to_descr}
+
+ # specializing the function according to the parameter description
+ new_scoped_function = (
+ self.scoped_functions[expr.function.name].with_descr(
+ combined_arg_id_to_descr))
+
+ # collecting the descriptors for args, kwargs, assignees
+ return set(((expr, new_scoped_function),))
+
+ def map_constant(self, expr):
+ return set()
+
+ map_variable = map_constant
+ map_function_symbol = map_constant
+
+def infer_arg_descr(kernel):
""" Specializes the kernel functions in way that the functions agree upon
shape and dimensions of the arguments too.
"""
- # The rest are to be hanfled by array calls. Which would need a mapper.
+ arg_description_modifier = ArgDescriptionInferer(kernel.scoped_functions)
+ pymbolic_calls_to_functions = set()
- new_insns = []
for insn in kernel.instructions:
+
+ if isinstance(insn, CallInstruction):
+ # In call instructions the assignees play an important in
+ # determining the arg_id_to_dtype
+ pymbolic_calls_to_functions.update(
+ arg_description_modifier(insn.expression,
+ assignees=insn.assignees))
if isinstance(insn, (MultiAssignmentBase, CInstruction)):
- expr = ArgDescriptionAdder(insn.expression)
- new_insns.append(insn.copy(expression=expr))
+ pymbolic_calls_to_functions.update(arg_description_modifier(
+ insn.expression))
elif isinstance(insn, _DataObliviousInstruction):
- new_insns.append()
+ pass
else:
raise NotImplementedError("arg_descr_inference for %s instruction" %
type(insn))
- # get the new scoped functions, in a similar fashion we did for type
- # inference
+ # making it the set of tuples a dict
+ pymbolic_calls_to_functions = dict(pymbolic_calls_to_functions)
+
+ # Now do the similar treatment as done for type inference.
+ from loopy.kernel.function_interface import (
+ register_pymbolic_calls_to_knl_callables)
+
+ return register_pymbolic_calls_to_knl_callables(kernel,
+ pymbolic_calls_to_functions)
- return kernel.copy(instructions=new_insns)
# }}}
@@ -2221,9 +2339,6 @@ def preprocess_kernel(kernel, device=None):
# Get them out of the way.
kernel = infer_unknown_types(kernel, expect_completion=False)
- print(kernel.instructions)
- print(kernel.scoped_functions)
- 1/0
# TODO: Specializng based on:
# 1. ArgDescriptors
@@ -2263,6 +2378,19 @@ def preprocess_kernel(kernel, device=None):
# have been established
kernel = check_atomic_loads(kernel)
+ kernel = infer_arg_descr(kernel)
+
+ print(75*'-')
+ print("This is after Type Inference")
+ for insn in kernel.instructions:
+ print(insn)
+ print(75*'-')
+ print('Linked Functions:')
+ for name, func in kernel.scoped_functions.items():
+ print(name, "=>", func)
+ print(75*'-')
+ 1/0
+
kernel = kernel.target.preprocess(kernel)
logger.info("%s: preprocess done" % kernel.name)
diff --git a/loopy/symbolic.py b/loopy/symbolic.py
index 23617c48b..8abda0f2a 100644
--- a/loopy/symbolic.py
+++ b/loopy/symbolic.py
@@ -723,19 +723,22 @@ class SubArrayRef(p.Expression):
starting_inames.append(iname)
return p.Subscript(self.subscript.aggregate, tuple(starting_inames))
- def get_inner_dim_tags(self, arg_dim_tags):
+ def get_sub_array_dim_tags_and_shape(self, arg_dim_tags, arg_shape):
""" Gives the dim tags for the inner inames.
This would be used for stride calculation in the child kernel.
This might need to go, once we start calculating the stride length
using the upper and lower bounds of the involved inames.
"""
from loopy.kernel.array import FixedStrideArrayDimTag as DimTag
- inner_dim_tags = []
- for dim_tag, iname in zip(arg_dim_tags, self.subscript.index_tuple):
+ sub_dim_tags = []
+ sub_shape = []
+ for dim_tag, axis_length, iname in zip(
+ arg_dim_tags, arg_shape, self.subscript.index_tuple):
if iname in self.swept_inames:
- inner_dim_tags.append(DimTag(dim_tag.stride))
+ sub_dim_tags.append(DimTag(dim_tag.stride))
+ sub_shape.append(axis_length)
- return inner_dim_tags
+ return sub_dim_tags, sub_shape
def __getinitargs__(self):
return (self.swept_inames, self.subscript)
diff --git a/loopy/type_inference.py b/loopy/type_inference.py
index 23aa379dd..bc8669528 100644
--- a/loopy/type_inference.py
+++ b/loopy/type_inference.py
@@ -25,9 +25,7 @@ THE SOFTWARE.
import six
from pymbolic.mapper import CombineMapper
-from loopy.symbolic import IdentityMapper, ScopedFunction
import numpy as np
-import re
from loopy.tools import is_integer
from loopy.types import NumpyType
@@ -36,9 +34,6 @@ from loopy.diagnostic import (
LoopyError,
TypeInferenceFailure, DependencyTypeInferenceFailure)
-from loopy.kernel.instruction import (MultiAssignmentBase, CInstruction,
- _DataObliviousInstruction)
-
import logging
logger = logging.getLogger(__name__)
@@ -515,59 +510,6 @@ class _DictUnionView:
raise KeyError(key)
-# {{{ duplicating the funciton name
-
-def next_indexed_name(name):
- FUNC_NAME = re.compile(r"^(?P<alpha>\S+?)_(?P<num>\d+?)$")
-
- match = FUNC_NAME.match(name)
-
- if match is None:
- if name[-1] == '_':
- return "{old_name}0".format(old_name=name)
- else:
- return "{old_name}_0".format(old_name=name)
-
- return "{alpha}_{num}".format(alpha=match.group('alpha'),
- num=int(match.group('num'))+1)
-
-# }}}
-
-
-# {{{ FunctionScopeChanger
-
-#TODO: Make it sophisticated
-
-class FunctionScopeChanger(IdentityMapper):
- def __init__(self, new_names):
- self.new_names = new_names
- self.new_names_set = frozenset(new_names.values())
-
- def map_call(self, expr):
- if expr in self.new_names:
- return type(expr)(
- ScopedFunction(self.new_names[expr]),
- tuple(self.rec(child)
- for child in expr.parameters))
- else:
- return IdentityMapper.map_call(self, expr)
-
- def map_call_with_kwargs(self, expr):
- if expr in self.new_names:
- return type(expr)(
- ScopedFunction(self.new_names[expr]),
- tuple(self.rec(child)
- for child in expr.parameters),
- dict(
- (key, self.rec(val))
- for key, val in six.iteritems(expr.kw_parameters))
- )
- else:
- return IdentityMapper.map_call_with_kwargs(self, expr)
-
-# }}}
-
-
# {{{ infer_unknown_types
def infer_unknown_types(kernel, expect_completion=False):
@@ -736,45 +678,11 @@ def infer_unknown_types(kernel, expect_completion=False):
args=[new_arg_dict[arg.name] for arg in kernel.args],
)
- # {{{ type specialization
-
- # TODO: These 2 dictionaries are inverse mapping of each other and help to keep
- # track of which ...(need to explain better)
- scoped_names_to_functions = pre_type_specialized_knl.scoped_functions
- scoped_functions_to_names = {}
- pymbolic_calls_to_new_names = {}
-
- for pymbolic_call, knl_callable in specialized_functions.items():
- if knl_callable not in scoped_functions_to_names:
- # need to make a new name deerived from the old name such that new
- # name in not present in new_scoped_name_to_function
- old_name = pymbolic_call.function.name
- new_name = next_indexed_name(old_name)
- while new_name in scoped_names_to_functions:
- new_name = next_indexed_name(new_name)
-
- scoped_names_to_functions[new_name] = knl_callable
- scoped_functions_to_names[knl_callable] = new_name
-
- pymbolic_calls_to_new_names[pymbolic_call] = (
- scoped_functions_to_names[knl_callable])
-
- # }}}
-
- new_insns = []
- scope_changer = FunctionScopeChanger(pymbolic_calls_to_new_names)
- for insn in pre_type_specialized_knl.instructions:
- if isinstance(insn, (MultiAssignmentBase, CInstruction)):
- expr = scope_changer(insn.expression)
- new_insns.append(insn.copy(expression=expr))
- elif isinstance(insn, _DataObliviousInstruction):
- new_insns.append(insn)
- else:
- raise NotImplementedError("Type Inference Specialization not"
- "implemented for %s instruciton" % type(insn))
+ from loopy.kernel.function_interface import (
+ register_pymbolic_calls_to_knl_callables)
+ return register_pymbolic_calls_to_knl_callables(
+ pre_type_specialized_knl, specialized_functions)
- return pre_type_specialized_knl.copy(scoped_functions=scoped_names_to_functions,
- instructions=new_insns)
# }}}
--
GitLab