From 31febb44daa2d6774a5a07258e3a3b900df78ffa Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 6 Jun 2016 21:48:34 -0500 Subject: [PATCH] Add for and with blocks to loopy-lang, for easier explicit iname decl, improve instruction parsing, adapt Rob's Bernstein tests --- loopy/kernel/creation.py | 275 ++++++++++++++++++++++++--------------- test/test_apps.py | 119 +++++++++-------- 2 files changed, 235 insertions(+), 159 deletions(-) diff --git a/loopy/kernel/creation.py b/loopy/kernel/creation.py index a1d8ece5b..0040a6d95 100644 --- a/loopy/kernel/creation.py +++ b/loopy/kernel/creation.py @@ -277,17 +277,23 @@ def parse_insn_options(opt_dict, options_str, assignee_names=None): # }}} -# {{{ parse instructions +# {{{ parse one instruction -OPTIONS_RE = re.compile( +WITH_OPTIONS_RE = re.compile( "^" - "(?P\s*)" + "\s*with\s*" "\{(?P.+)\}" "\s*$") +FOR_RE = re.compile( + "^" + "\s*for\s*" + "(?P[ ,\w]+)" + "\s*$") + INSN_RE = re.compile( "^" - "(?P\s*)" + "\s*" "(?P.+?)" "\s*(?.+?)" @@ -298,7 +304,7 @@ SUBST_RE = re.compile( r"^\s*(?P.+?)\s*:=\s*(?P.+)\s*$") -def parse_insn(insn): +def parse_insn(groups, insn_options): """ :return: a tuple ``(insn, inames_to_dup)``, where insn is a :class:`Assignment`, a :class:`CallInstruction`, @@ -308,18 +314,6 @@ def parse_insn(insn): import loopy as lp - insn_match = INSN_RE.match(insn) - subst_match = SUBST_RE.match(insn) - if insn_match is not None and subst_match is not None: - raise RuntimeError("instruction parse error: %s" % insn) - - if insn_match is not None: - groups = insn_match.groupdict() - elif subst_match is not None: - groups = subst_match.groupdict() - else: - raise RuntimeError("instruction parse error at '%s'" % insn) - from loopy.symbolic import parse try: lhs = parse(groups["lhs"]) @@ -335,40 +329,9 @@ def parse_insn(insn): "the following error occurred:" % groups["rhs"]) raise - from pymbolic.primitives import Variable, Call, Subscript + from pymbolic.primitives import Variable, Subscript from loopy.symbolic import TypeAnnotation - # {{{ deal with subst rules - - if subst_match is not None: - assert insn_match is None - if isinstance(lhs, Variable): - subst_name = lhs.name - arg_names = [] - elif isinstance(lhs, Call): - if not isinstance(lhs.function, Variable): - raise RuntimeError("Invalid substitution rule left-hand side") - subst_name = lhs.function.name - arg_names = [] - - for i, arg in enumerate(lhs.parameters): - if not isinstance(arg, Variable): - raise RuntimeError("Invalid substitution rule " - "left-hand side: %s--arg number %d " - "is not a variable" % (lhs, i)) - arg_names.append(arg.name) - else: - raise RuntimeError("Invalid substitution rule left-hand side") - - return SubstitutionRule( - name=subst_name, - arguments=tuple(arg_names), - expression=rhs), [] - - assert insn_match is not None - - # }}} - if not isinstance(lhs, tuple): lhs = (lhs,) @@ -403,12 +366,12 @@ def parse_insn(insn): del new_lhs insn_options = parse_insn_options( - get_default_insn_options_dict(), + insn_options, groups["options"], assignee_names=assignee_names) insn_id = insn_options.pop("insn_id", None) - inames_to_dup = insn_options.pop("inames_to_dup", None) + inames_to_dup = insn_options.pop("inames_to_dup", []) kwargs = dict( id=( @@ -422,37 +385,159 @@ def parse_insn(insn): lhs, rhs, temp_var_types, **kwargs ), inames_to_dup +# }}} + + +# {{{ parse_subst_rule + +def parse_subst_rule(groups): + from loopy.symbolic import parse + try: + lhs = parse(groups["lhs"]) + except: + print("While parsing left hand side '%s', " + "the following error occurred:" % groups["lhs"]) + raise + + try: + rhs = parse(groups["rhs"]) + except: + print("While parsing right hand side '%s', " + "the following error occurred:" % groups["rhs"]) + raise + + from pymbolic.primitives import Variable, Call + if isinstance(lhs, Variable): + subst_name = lhs.name + arg_names = [] + elif isinstance(lhs, Call): + if not isinstance(lhs.function, Variable): + raise RuntimeError("Invalid substitution rule left-hand side") + subst_name = lhs.function.name + arg_names = [] + + for i, arg in enumerate(lhs.parameters): + if not isinstance(arg, Variable): + raise RuntimeError("Invalid substitution rule " + "left-hand side: %s--arg number %d " + "is not a variable" % (lhs, i)) + arg_names.append(arg.name) + else: + raise RuntimeError("Invalid substitution rule left-hand side") + + return SubstitutionRule( + name=subst_name, + arguments=tuple(arg_names), + expression=rhs) + +# }}} + + +# {{{ parse_instructions + +def parse_instructions(instructions, defines): + if isinstance(instructions, str): + instructions = [instructions] + + parsed_instructions = [] + + # {{{ pass 1: interning, comments, whitespace, defines -def parse_if_necessary(insn, defines): - if isinstance(insn, InstructionBase): - yield insn.copy( - id=intern(insn.id) if isinstance(insn.id, str) else insn.id, - depends_on=frozenset(intern(dep) for dep in insn.depends_on), - groups=frozenset(intern(grp) for grp in insn.groups), - conflicts_with_groups=frozenset( - intern(grp) for grp in insn.conflicts_with_groups), - forced_iname_deps=frozenset( - intern(iname) for iname in insn.forced_iname_deps), - predicates=frozenset( - intern(pred) for pred in insn.predicates), - ), [] - return - elif not isinstance(insn, str): - raise TypeError("Instructions must be either an Instruction " - "instance or a parseable string. got '%s' instead." - % type(insn)) - - for insn in insn.split("\n"): - comment_start = insn.find("#") - if comment_start >= 0: - insn = insn[:comment_start] - - insn = insn.strip() - if not insn: + for insn in instructions: + if isinstance(insn, InstructionBase): + parsed_instructions.append( + insn.copy( + id=intern(insn.id) if isinstance(insn.id, str) else insn.id, + depends_on=frozenset(intern(dep) for dep in insn.depends_on), + groups=frozenset(intern(grp) for grp in insn.groups), + conflicts_with_groups=frozenset( + intern(grp) for grp in insn.conflicts_with_groups), + forced_iname_deps=frozenset( + intern(iname) for iname in insn.forced_iname_deps), + predicates=frozenset( + intern(pred) for pred in insn.predicates), + )) + continue + + elif not isinstance(insn, str): + raise TypeError("Instructions must be either an Instruction " + "instance or a parseable string. got '%s' instead." + % type(insn)) + + for insn in insn.split("\n"): + comment_start = insn.find("#") + if comment_start >= 0: + insn = insn[:comment_start] + + insn = insn.strip() + if not insn: + continue + + for sub_insn in expand_defines(insn, defines, single_valued=False): + parsed_instructions.append(sub_insn) + + # }}} + + instructions = parsed_instructions + parsed_instructions = [] + substitutions = {} + inames_to_dup = [] # one for each parsed_instruction + + # {{{ pass 2: parsing + + insn_options_stack = [get_default_insn_options_dict()] + + for insn in instructions: + if isinstance(insn, InstructionBase): + parsed_instructions.append(insn) + inames_to_dup.append([]) + continue + + with_options_match = WITH_OPTIONS_RE.match(insn) + if with_options_match is not None: + insn_options_stack.append( + parse_insn_options( + insn_options_stack[-1], + with_options_match.group("options"))) + continue + + for_match = FOR_RE.match(insn) + if for_match is not None: + options = insn_options_stack[-1].copy() + options["forced_iname_deps"] = ( + options.get("forced_iname_deps", frozenset()) + | frozenset( + iname.strip() + for iname in for_match.group("inames").split(","))) + options["forced_iname_deps_is_final"] = True + + insn_options_stack.append(options) + del options + continue + + if insn == "end": + insn_options_stack.pop() + continue + + subst_match = SUBST_RE.match(insn) + if subst_match is not None: + subst = parse_subst_rule(subst_match.groupdict()) + substitutions[subst.name] = subst + continue + + insn_match = INSN_RE.match(insn) + if insn_match is not None: + insn, insn_inames_to_dup = parse_insn( + insn_match.groupdict(), insn_options_stack[-1]) + parsed_instructions.append(insn) + inames_to_dup.append(insn_inames_to_dup) continue - for sub_insn in expand_defines(insn, defines, single_valued=False): - yield parse_insn(sub_insn) + raise RuntimeError("instruction parse error: %s" % insn) + + # }}} + + return parsed_instructions, inames_to_dup, substitutions # }}} @@ -1275,34 +1360,8 @@ def make_kernel(domains, instructions, kernel_data=["..."], **kwargs): # }}} - # {{{ instruction/subst parsing - - parsed_instructions = [] - kwargs["substitutions"] = substitutions = {} - inames_to_dup = [] - - if isinstance(instructions, str): - instructions = [instructions] - - for insn in instructions: - for new_insn, insn_inames_to_dup in parse_if_necessary(insn, defines): - if isinstance(new_insn, InstructionBase): - parsed_instructions.append(new_insn) - - # Need to maintain 1-to-1 correspondence to instructions - inames_to_dup.append(insn_inames_to_dup) - - elif isinstance(new_insn, SubstitutionRule): - substitutions[new_insn.name] = new_insn - - assert not insn_inames_to_dup - else: - raise RuntimeError("unexpected type in instruction parsing") - - instructions = parsed_instructions - del parsed_instructions - - # }}} + instructions, inames_to_dup, substitutions = \ + parse_instructions(instructions, defines) # {{{ find/create isl_context @@ -1326,6 +1385,8 @@ def make_kernel(domains, instructions, kernel_data=["..."], **kwargs): kernel_args = arg_guesser.convert_names_to_full_args(kernel_args) kernel_args = arg_guesser.guess_kernel_args_if_requested(kernel_args) + kwargs["substitutions"] = substitutions + from loopy.kernel import LoopKernel knl = LoopKernel(domains, instructions, kernel_args, temporary_variables=temporary_variables, diff --git a/test/test_apps.py b/test/test_apps.py index a73356671..6dfc6dc0c 100644 --- a/test/test_apps.py +++ b/test/test_apps.py @@ -183,22 +183,26 @@ def test_rob_stroud_bernstein(ctx_factory): 0 <= i2 < nqp1d and \ 0 <= alpha1 <= deg and 0 <= alpha2 <= deg-alpha1 }", """ - <> xi = qpts[1, i2] {inames=+el} + for el,i2 + <> xi = qpts[1, i2] <> s = 1-xi <> r = xi/s - <> aind = 0 {id=aind_init,inames=+i2:el} - - <> w = s**(deg-alpha1) {id=init_w} - - tmp[el,alpha1,i2] = tmp[el,alpha1,i2] + w * coeffs[aind] \ - {id=write_tmp,inames=+alpha2} - w = w * r * ( deg - alpha1 - alpha2 ) / (1 + alpha2) \ - {id=update_w,dep=init_w:write_tmp} - aind = aind + 1 \ - {id=aind_incr,\ - dep=aind_init:write_tmp:update_w, \ - inames=+el:i2:alpha1:alpha2} - """, + <> aind = 0 {id=aind_init} + + for alpha1 + <> w = s**(deg-alpha1) {id=init_w} + + for alpha2 + tmp[el,alpha1,i2] = tmp[el,alpha1,i2] + w * coeffs[aind] \ + {id=write_tmp} + w = w * r * ( deg - alpha1 - alpha2 ) / (1 + alpha2) \ + {id=update_w,dep=init_w:write_tmp} + aind = aind + 1 \ + {id=aind_incr,dep=aind_init:write_tmp:update_w} + end + end + end + """, [ # Must declare coeffs to have "no" shape, to keep loopy # from trying to figure it out the shape automatically. @@ -230,52 +234,63 @@ def test_rob_stroud_bernstein_full(ctx_factory): # NOTE: result would have to be zero-filled beforehand knl = lp.make_kernel( - "{[el, i2, alpha1,alpha2, i1_2, alpha1_2, i2_2]: \ - 0 <= el < nels and \ - 0 <= i2 < nqp1d and \ - 0 <= alpha1 <= deg and 0 <= alpha2 <= deg-alpha1 and\ - \ - 0 <= i1_2 < nqp1d and \ - 0 <= alpha1_2 <= deg and \ - 0 <= i2_2 < nqp1d \ - }", - """ - <> xi = qpts[1, i2] {inames=+el} + "{[el, i2, alpha1,alpha2, i1_2, alpha1_2, i2_2]: \ + 0 <= el < nels and \ + 0 <= i2 < nqp1d and \ + 0 <= alpha1 <= deg and 0 <= alpha2 <= deg-alpha1 and\ + \ + 0 <= i1_2 < nqp1d and \ + 0 <= alpha1_2 <= deg and \ + 0 <= i2_2 < nqp1d \ + }", + """ + for el + for i2 + <> xi = qpts[1, i2] <> s = 1-xi <> r = xi/s - <> aind = 0 {id=aind_init,inames=+i2:el} - - <> w = s**(deg-alpha1) {id=init_w} - - <> tmp[alpha1,i2] = tmp[alpha1,i2] + w * coeffs[aind] \ - {id=write_tmp,inames=+alpha2} - w = w * r * ( deg - alpha1 - alpha2 ) / (1 + alpha2) \ - {id=update_w,dep=init_w:write_tmp} - aind = aind + 1 \ - {id=aind_incr,\ - dep=aind_init:write_tmp:update_w, \ - inames=+el:i2:alpha1:alpha2} - - <> xi2 = qpts[0, i1_2] {dep=aind_incr,inames=+el} + <> aind = 0 {id=aind_init} + + for alpha1 + <> w = s**(deg-alpha1) {id=init_w} + + <> tmp[alpha1,i2] = tmp[alpha1,i2] + w * coeffs[aind] \ + {id=write_tmp} + for alpha2 + w = w * r * ( deg - alpha1 - alpha2 ) / (1 + alpha2) \ + {id=update_w,dep=init_w:write_tmp} + aind = aind + 1 \ + {id=aind_incr,dep=aind_init:write_tmp:update_w} + end + end + end + + for i1_2 + <> xi2 = qpts[0, i1_2] {dep=aind_incr} <> s2 = 1-xi2 <> r2 = xi2/s2 <> w2 = s2**deg - result[el, i1_2, i2_2] = result[el, i1_2, i2_2] + \ - w2 * tmp[alpha1_2, i2_2] \ - {inames=el:alpha1_2:i1_2:i2_2} + for alpha1_2 + for i2_2 + result[el, i1_2, i2_2] = result[el, i1_2, i2_2] + \ + w2 * tmp[alpha1_2, i2_2] + end - w2 = w2 * r2 * (deg-alpha1_2) / (1+alpha1_2) - """, - [ - # Must declare coeffs to have "no" shape, to keep loopy - # from trying to figure it out the shape automatically. + w2 = w2 * r2 * (deg-alpha1_2) / (1+alpha1_2) + end + end + end + """, + [ + # Must declare coeffs to have "no" shape, to keep loopy + # from trying to figure it out the shape automatically. - lp.GlobalArg("coeffs", None, shape=None), - "..." - ], - assumptions="deg>=0 and nels>=1" - ) + lp.GlobalArg("coeffs", None, shape=None), + "..." + ], + assumptions="deg>=0 and nels>=1" + ) knl = lp.fix_parameters(knl, nqp1d=7, deg=4) -- GitLab