From 111a4f1da965905e8a5166a9de97e11aa6c70ebd Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner <inform@tiker.net> Date: Sun, 28 Aug 2011 21:01:38 +0200 Subject: [PATCH] Lots of post-CFEM-quadrature changes. Implement user control over slab decomposition. Use CSEs instead of register prefetch. Remove CodeGenerationState. Carry CSE replacement info in CodeGenerationMapper. Carry CCM in ExecutionDomain. Allow user to specify (unchecked) assumptions on parameters, to avoid unnecessary conditionals. Parse desired operation from strings, extract CSEs. Support for equality constraints in for loops. --- examples/quadrature.py | 23 ++++---- loopy/__init__.py | 9 ++- loopy/codegen/__init__.py | 65 ++++++++++----------- loopy/codegen/bounds.py | 92 +++++++++++++++++------------ loopy/codegen/dispatch.py | 53 ++++++----------- loopy/codegen/loop_dim.py | 111 +++++++++++------------------------ loopy/codegen/prefetch.py | 17 +++--- loopy/compiled.py | 2 +- loopy/kernel.py | 120 ++++++++++++++++++++++++++------------ loopy/prefetch.py | 6 +- loopy/symbolic.py | 63 +++++++++++++++++--- test/test_matmul.py | 71 +++++++++++----------- 12 files changed, 338 insertions(+), 294 deletions(-) diff --git a/examples/quadrature.py b/examples/quadrature.py index 9bf4e6f21..c39d6c99a 100644 --- a/examples/quadrature.py +++ b/examples/quadrature.py @@ -34,17 +34,14 @@ def build_mass_mat_maker(ctx_factory=cl.create_some_context): Nb = 3 Nv = 3 Nq = 3*3 - Nc = 1600 - from pymbolic import var - m, w, det_j, phi, c, i, j, q = [var(s) for s in "m w det_j phi c i j q".split()] knl = lp.LoopKernel(ctx.devices[0], "[ncells] -> {[c,i,j,q]: 0<=c<ncells and 0 <= i < %(Nv)s " "and 0<=j<%(Nb)s and 0<=q<%(Nq)s}" % dict( Nv=Nv, Nb=Nb, Nq=Nq), [ - (m[c, i, j], w[q]*det_j[c]*phi[i,q]*phi[j,q]) + "m[c,i,j] = w[q]*det_j[c]*phi[i,q]*phi[j,q]", ], [ lp.ArrayArg("m", dtype, shape=(Nc, Nv, Nb)), @@ -54,17 +51,16 @@ def build_mass_mat_maker(ctx_factory=cl.create_some_context): lp.ScalarArg("ncells", np.int32, approximately=1000), ], name="mass_mat", - iname_to_tag=dict(i="l.0", j="l.1") + iname_to_tag=dict(i="l.0", j="l.1"), + assumptions="ncells >= 1" ) - knl = lp.split_dimension(knl, "c", 8, outer_tag="g.0", inner_tag="l.2") - knl = lp.add_prefetch(knl, "det_j", ["c_inner"]) + knl = lp.split_dimension(knl, "c", 8, inner_tag="l.2", + outer_slab_increments=(0,0)) + knl = lp.split_dimension(knl, "c_outer", 8, outer_tag="g.0", + outer_slab_increments=(0,0)) # fix reg prefetch - # fix redundant slab generation - - # FIXME - #knl = lp.split_dimension(knl, "c", 8, inner_tag="l.2") - #knl = lp.split_dimension(knl, "c_outer", 8, outer_tag="g.0") + knl = lp.add_prefetch(knl, "det_j", ["c_inner"]) #ilp = 4 #knl = lp.split_dimension(knl, "i", 2, outer_tag="g.0", inner_tag="l.1") @@ -77,7 +73,8 @@ def build_mass_mat_maker(ctx_factory=cl.create_some_context): #knl = lp.add_prefetch(knl, 'b', ["j_inner_outer", "j_inner_inner", "k_inner"]) #assert knl.get_problems({})[0] <= 2 - kernel_gen = (lp.insert_register_prefetches(knl) + kernel_gen = (knl + #kernel_gen = (lp.insert_register_prefetches(knl) for knl in lp.generate_loop_schedules(knl)) if False: diff --git a/loopy/__init__.py b/loopy/__init__.py index 147efabf9..46bc404e5 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -21,24 +21,23 @@ register_mpz_with_pymbolic() # TODO: Custom reductions per red. axis # TODO: Functions # TODO: Common subexpressions -# TODO: Parse ops from string +# TODO: Array common subexpressions # FIXME: support non-reductive dimensions (what did I mean here?) # FIXME: write names should be assigned during scheduling +# FIXME: screwy lower bounds in ILP # TODO: Divisibility # TODO: Try, fix indirect addressing -# TODO: More user control for slab opt # TODO: Implement GT200 matmul, Fermi matmul, DG # TODO: DMA engine threads? -# TODO: Specify initial implemented domain. -# (to filter away unnecessary conditions on parameters) # TODO: Deal with equalities that crop up. +# TODO: Better user feedback. # Later: # ------ # TODO: Try different kernels -# TODO: - Tricky: Convolution, FD +# TODO: - Tricky: Convolution, Stencil # TODO: Separate all-bulk from non-bulk kernels. (maybe?) (#ifdef?) # TODO: implement efficient ceil_div? (as opposed to floor_div) # TODO: why are corner cases inefficient? diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index e8f10a1a8..6b9cb966f 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -2,7 +2,6 @@ from __future__ import division from pytools import Record import numpy as np -from pymbolic.mapper.stringifier import PREC_NONE @@ -95,50 +94,56 @@ def add_comment(cmt, code): # {{{ main code generation entrypoint -class CodeGenerationState(Record): - __slots__ = ["c_code_mapper", "try_slab_partition"] +class ExecutionSubdomain(Record): + __slots__ = ["implemented_domain", "c_code_mapper"] + + def __init__(self, implemented_domain, c_code_mapper): + Record.__init__(self, + implemented_domain=implemented_domain, + c_code_mapper=c_code_mapper) + + def intersect(self, set): + return ExecutionSubdomain( + self.implemented_domain.intersect(set), + self.c_code_mapper) class ExecutionDomain(object): - def __init__(self, implemented_domain, assignments_and_impl_domains=None): + def __init__(self, implemented_domain, c_code_mapper, subdomains=None): """ :param implemented_domain: The entire implemented domain, i.e. all constraints that have been enforced so far. - :param assignments_and_impl_domains: a list of tuples - (assignments, implemented_domain), where *assignments* - is a list of :class:`cgen.Initializer` instances - and *implemented_domain* is the implemented domain to which - the situation produced by the assignments corresponds. + :param subdomains: a list of :class:`ExecutionSubdomain` + instances. - The point of this being is a list is the implementation of + The point of this being a list is the implementation of ILP, and each entry represents a 'fake-parallel' trip through the - ILP'd loop. + ILP'd loop, with the requisite implemented_domain + and a C code mapper that realizes the necessary assignments. + :param c_code_mapper: A C code mapper that does not take per-ILP + assignments into account. """ - if assignments_and_impl_domains is None: - assignments_and_impl_domains = [([], implemented_domain)] self.implemented_domain = implemented_domain - self.assignments_and_impl_domains = assignments_and_impl_domains - - def __len__(self): - return len(self.assignments_and_impl_domains) + if subdomains is None: + self.subdomains = [ + ExecutionSubdomain(implemented_domain, c_code_mapper)] + else: + self.subdomains = subdomains - def __iter__(self): - return iter(self.assignments_and_impl_domains) + self.c_code_mapper = c_code_mapper def intersect(self, set): return ExecutionDomain( self.implemented_domain.intersect(set), - [(assignments, implemented_domain.intersect(set)) - for assignments, implemented_domain - in self.assignments_and_impl_domains]) + self.c_code_mapper, + [subd.intersect(set) for subd in self.subdomains]) def get_the_one_domain(self): - assert len(self.assignments_and_impl_domains) == 1 + assert len(self.subdomains) == 1 return self.implemented_domain - def generate_code(kernel): from loopy.codegen.prefetch import preprocess_prefetch kernel = preprocess_prefetch(kernel) @@ -151,10 +156,7 @@ def generate_code(kernel): CLLocal, CLImage, CLConstant) from loopy.symbolic import LoopyCCodeMapper - my_ccm = LoopyCCodeMapper(kernel) - - def ccm(expr, prec=PREC_NONE): - return my_ccm(expr, prec) + ccm = LoopyCCodeMapper(kernel) # {{{ build top-level @@ -266,13 +268,10 @@ def generate_code(kernel): # }}} - cgs = CodeGenerationState(c_code_mapper=ccm, try_slab_partition=True) - from loopy.codegen.dispatch import build_loop_nest - import islpy as isl - gen_code = build_loop_nest(cgs, kernel, 0, - ExecutionDomain(isl.Set.universe(kernel.space))) + gen_code = build_loop_nest(kernel, 0, + ExecutionDomain( kernel.assumptions, c_code_mapper=ccm)) body.extend([Line(), gen_code.ast]) #print "# conditionals: %d" % gen_code.num_conditionals diff --git a/loopy/codegen/bounds.py b/loopy/codegen/bounds.py index e602e951e..4c509d4c7 100644 --- a/loopy/codegen/bounds.py +++ b/loopy/codegen/bounds.py @@ -2,7 +2,7 @@ from __future__ import division import islpy as isl from islpy import dim_type -from pymbolic.mapper.stringifier import PREC_NONE +import numpy as np @@ -192,45 +192,61 @@ def wrap_in_for_from_constraints(ccm, iname, constraint_bset, stmt): cfm = CommutativeConstantFoldingMapper() from loopy.symbolic import constraint_to_expr - from pytools import any - if any(cns.is_equality() for cns in constraints): - raise NotImplementedError("equality constraints for 'for' loops") - else: - start_exprs = [] - end_conds = [] - - for cns in constraints: - rhs, iname_coeff = constraint_to_expr(cns, except_name=iname) - - if iname_coeff == 0: - continue - elif iname_coeff < 0: - from pymbolic import var - rhs += iname_coeff*var(iname) - end_conds.append("%s >= 0" % - ccm(cfm(expand(rhs)))) - else: # iname_coeff > 0 - kind, bound = solve_constraint_for_bound(cns, iname) - assert kind == ">=" - start_exprs.append(bound) - - if len(start_exprs) > 1: - from pymbolic.primitives import Max - start_expr = Max(start_exprs) - elif len(start_exprs) == 1: - start_expr, = start_exprs - else: - raise RuntimeError("no starting value found for 'for' loop in '%s'" - % iname) + start_exprs = [] + end_conds = [] + equality_exprs = [] - from cgen import For - from loopy.codegen import wrap_in - return wrap_in(For, - "int %s = %s" % (iname, ccm(start_expr, PREC_NONE)), - " && ".join(end_conds), - "++%s" % iname, - stmt) + for cns in constraints: + rhs, iname_coeff = constraint_to_expr(cns, except_name=iname) + + if iname_coeff == 0: + continue + + if cns.is_equality(): + kind, bound = solve_constraint_for_bound(cns, iname) + assert kind == "==" + equality_exprs.append(bound) + elif iname_coeff < 0: + from pymbolic import var + rhs += iname_coeff*var(iname) + end_conds.append("%s >= 0" % + ccm(cfm(expand(rhs)))) + else: # iname_coeff > 0 + kind, bound = solve_constraint_for_bound(cns, iname) + assert kind == ">=" + start_exprs.append(bound) + + if equality_exprs: + assert len(equality_exprs) == 1 + + equality_expr, = equality_exprs + + from loopy.codegen import gen_code_block + from cgen import Initializer, POD, Const, Line + return gen_code_block([ + Initializer(Const(POD(np.int32, iname)), + ccm(equality_expr)), + Line(), + stmt, + ]) + else: + if len(start_exprs) > 1: + from pymbolic.primitives import Max + start_expr = Max(start_exprs) + elif len(start_exprs) == 1: + start_expr, = start_exprs + else: + raise RuntimeError("no starting value found for 'for' loop in '%s'" + % iname) + + from cgen import For + from loopy.codegen import wrap_in + return wrap_in(For, + "int %s = %s" % (iname, ccm(start_expr)), + " && ".join(end_conds), + "++%s" % iname, + stmt) # }}} diff --git a/loopy/codegen/dispatch.py b/loopy/codegen/dispatch.py index f13b501a5..e1f688a05 100644 --- a/loopy/codegen/dispatch.py +++ b/loopy/codegen/dispatch.py @@ -6,10 +6,10 @@ from loopy.codegen import ExecutionDomain, gen_code_block -def build_loop_nest(cgs, kernel, sched_index, exec_domain, no_conditional_check=False): +def build_loop_nest(kernel, sched_index, exec_domain, no_conditional_check=False): assert isinstance(exec_domain, ExecutionDomain) - ccm = cgs.c_code_mapper + ccm = exec_domain.c_code_mapper from cgen import (POD, Initializer, Assign, Statement as S, Line) @@ -37,7 +37,7 @@ def build_loop_nest(cgs, kernel, sched_index, exec_domain, no_conditional_check= exec_domain = exec_domain.intersect(exec_domain_restriction) - inner = build_loop_nest(cgs, kernel, sched_index, exec_domain, + inner = build_loop_nest(kernel, sched_index, exec_domain, no_conditional_check=True) from loopy.codegen import wrap_in_if @@ -57,30 +57,21 @@ def build_loop_nest(cgs, kernel, sched_index, exec_domain, no_conditional_check= valid_index_vars = get_valid_check_vars(kernel, sched_index, allow_ilp=True) bounds_check_lists = [ - generate_bounds_checks_code(ccm, kernel.domain, - valid_index_vars, impl_domain) - for assignments, impl_domain in - exec_domain] + generate_bounds_checks_code(subd.c_code_mapper, kernel.domain, + valid_index_vars, subd.implemented_domain) + for subd in exec_domain.subdomains] result = [] for lvalue, expr in kernel.instructions: - for i, (assignments, impl_domain) in \ - enumerate(exec_domain): - - my_block = [] - if assignments: - my_block.extend(assignments+[Line()]) - + for i, subd in enumerate(exec_domain.subdomains): assert isinstance(lvalue, Subscript) name = lvalue.aggregate.name from loopy.codegen import wrap_in_if - my_block.append( - wrap_in_if( + result.append(wrap_in_if( bounds_check_lists[i], S("tmp_%s_%d += %s" - % (name, i, ccm(expr))))) - result.append(gen_code_block(my_block)) + % (name, i, subd.c_code_mapper(expr))))) return gen_code_block(result) @@ -109,42 +100,36 @@ def build_loop_nest(cgs, kernel, sched_index, exec_domain, no_conditional_check= else: func = generate_sequential_loop_dim_code - return func(cgs, kernel, sched_index, exec_domain) + return func(kernel, sched_index, exec_domain) elif isinstance(sched_item, WriteOutput): result = ( [Initializer(POD(kernel.arg_dict[lvalue.aggregate.name].dtype, "tmp_%s_%d" % (lvalue.aggregate.name, i)), 0) - for i in range(len(exec_domain)) + for i in range(len(exec_domain.subdomains)) for lvalue, expr in kernel.instructions] +[Line()] - +[build_loop_nest(cgs, kernel, sched_index+1, + +[build_loop_nest(kernel, sched_index+1, exec_domain)] +[Line()]) - for i, (idx_assignments, impl_domain) in \ - enumerate(exec_domain): + for i, subd in enumerate(exec_domain.subdomains): for lvalue, expr in kernel.instructions: - assignment = Assign(ccm(lvalue), "tmp_%s_%d" % ( + assignment = Assign(subd.c_code_mapper(lvalue), "tmp_%s_%d" % ( lvalue.aggregate.name, i)) wrapped_assign = wrap_in_bounds_checks( - ccm, kernel.domain, + subd.c_code_mapper, kernel.domain, get_valid_check_vars(kernel, sched_index, allow_ilp=True), - impl_domain, assignment) - - cb = [] - if idx_assignments: - cb.extend(idx_assignments+[Line()]) - cb.append(wrapped_assign) + subd.implemented_domain, assignment) - result.append(gen_code_block(cb)) + result.append(wrapped_assign) return gen_code_block(result) elif isinstance(sched_item, LocalMemoryPrefetch): from loopy.codegen.prefetch import generate_prefetch_code - return generate_prefetch_code(cgs, kernel, sched_index, + return generate_prefetch_code(kernel, sched_index, exec_domain) elif isinstance(sched_item, RegisterPrefetch): @@ -159,7 +144,7 @@ def build_loop_nest(cgs, kernel, sched_index, exec_domain, no_conditional_check= % (agg_name, ccm(sched_item.subscript_expr.index)))), - build_loop_nest(cgs, kernel, sched_index+1, exec_domain) + build_loop_nest(kernel, sched_index+1, exec_domain) ]) else: diff --git a/loopy/codegen/loop_dim.py b/loopy/codegen/loop_dim.py index 7c799a84c..888682e5a 100644 --- a/loopy/codegen/loop_dim.py +++ b/loopy/codegen/loop_dim.py @@ -32,86 +32,46 @@ def get_simple_loop_bounds(kernel, sched_index, iname, implemented_domain): # {{{ conditional-minimizing slab decomposition -def get_slab_decomposition(cgs, kernel, sched_index, exec_domain): +def get_slab_decomposition(kernel, sched_index, exec_domain): from loopy.isl import block_shift_constraint, negate_constraint - ccm = cgs.c_code_mapper + ccm = exec_domain.c_code_mapper space = kernel.space iname = kernel.schedule[sched_index].iname tag = kernel.iname_to_tag.get(iname) - # {{{ attempt slab partition to reduce conditional count - lb_cns_orig, ub_cns_orig = get_simple_loop_bounds(kernel, sched_index, iname, exec_domain.implemented_domain) - # jostle the constant in {lb,ub}_cns to see if we can get - # fewer conditionals in the bulk middle segment - - class TrialRecord(Record): - pass - - if (cgs.try_slab_partition - and "outer" in iname): - trial_cgs = cgs.copy(try_slab_partition=False) - trials = [] - - for lower_incr, upper_incr in [ (0,0), (0,-1), ]: - - lb_cns = block_shift_constraint(lb_cns_orig, iname, -lower_incr) - ub_cns = block_shift_constraint(ub_cns_orig, iname, -upper_incr) - - bulk_slab = (isl.Set.universe(kernel.space) - .add_constraint(lb_cns) - .add_constraint(ub_cns)) - bulk_exec_domain = exec_domain.intersect(bulk_slab) - inner = build_loop_nest(trial_cgs, kernel, sched_index+1, - bulk_exec_domain) - - trials.append((TrialRecord( - lower_incr=lower_incr, - upper_incr=upper_incr, - bulk_slab=bulk_slab), - (inner.num_conditionals, - # when all num_conditionals are equal, choose the - # one with the smallest bounds changes - abs(upper_incr)+abs(lower_incr)))) - - from pytools import argmin2 - chosen = argmin2(trials) - else: - bulk_slab = (isl.Set.universe(kernel.space) - .add_constraint(lb_cns_orig) - .add_constraint(ub_cns_orig)) - chosen = TrialRecord( - lower_incr=0, - upper_incr=0, - bulk_slab=bulk_slab) - - # }}} + lower_incr, upper_incr = kernel.iname_slab_increments.get(iname, (0, 0)) # {{{ build slabs slabs = [] - if chosen.lower_incr: + if lower_incr: slabs.append(("initial", isl.Set.universe(kernel.space) .add_constraint(lb_cns_orig) .add_constraint(ub_cns_orig) .add_constraint( negate_constraint( block_shift_constraint( - lb_cns_orig, iname, -chosen.lower_incr))))) + lb_cns_orig, iname, -lower_incr))))) - slabs.append(("bulk", chosen.bulk_slab)) + slabs.append(("bulk", + (isl.Set.universe(kernel.space) + .add_constraint( + block_shift_constraint(lb_cns_orig, iname, -lower_incr)) + .add_constraint( + block_shift_constraint(ub_cns_orig, iname, -upper_incr))))) - if chosen.upper_incr: + if upper_incr: slabs.append(("final", isl.Set.universe(kernel.space) .add_constraint(ub_cns_orig) .add_constraint(lb_cns_orig) .add_constraint( negate_constraint( block_shift_constraint( - ub_cns_orig, iname, -chosen.upper_incr))))) + ub_cns_orig, iname, -upper_incr))))) # }}} @@ -121,13 +81,13 @@ def get_slab_decomposition(cgs, kernel, sched_index, exec_domain): # {{{ unrolled/ILP loops -def generate_unroll_or_ilp_code(cgs, kernel, sched_index, exec_domain): +def generate_unroll_or_ilp_code(kernel, sched_index, exec_domain): from loopy.isl import block_shift_constraint from loopy.codegen.bounds import solve_constraint_for_bound from cgen import (POD, Assign, Line, Statement as S, Initializer, Const) - ccm = cgs.c_code_mapper + ccm = exec_domain.c_code_mapper space = kernel.space iname = kernel.schedule[sched_index].iname tag = kernel.iname_to_tag.get(iname) @@ -162,8 +122,7 @@ def generate_unroll_or_ilp_code(cgs, kernel, sched_index, exec_domain): for i, slab in generate_idx_eq_slabs(): new_exec_domain = exec_domain.intersect(slab) - inner = build_loop_nest(cgs, kernel, sched_index+1, - new_exec_domain) + inner = build_loop_nest(kernel, sched_index+1, new_exec_domain) if isinstance(tag, TAG_UNROLL_STATIC): result.extend([ @@ -175,24 +134,25 @@ def generate_unroll_or_ilp_code(cgs, kernel, sched_index, exec_domain): return gen_code_block(result) elif isinstance(tag, TAG_ILP): - new_aaid = [] - for assignments, implemented_domain in exec_domain: + new_subdomains = [] + for subd in exec_domain.subdomains: for i, single_slab in generate_idx_eq_slabs(): - assignments = assignments + [ - Initializer(Const(POD(np.int32, iname)), ccm(lower_bound+i))] - new_aaid.append((assignments, - implemented_domain.intersect(single_slab))) - - assignments = [] + from loopy.codegen import ExecutionSubdomain + new_subdomains.append( + ExecutionSubdomain( + subd.implemented_domain.intersect(single_slab), + subd.c_code_mapper.copy_and_assign( + iname, lower_bound+i))) overall_slab = (isl.Set.universe(kernel.space) .add_constraint(lower_cns) .add_constraint(upper_cns)) - return build_loop_nest(cgs, kernel, sched_index+1, + return build_loop_nest(kernel, sched_index+1, ExecutionDomain( exec_domain.implemented_domain.intersect(overall_slab), - new_aaid)) + exec_domain.c_code_mapper, + new_subdomains)) else: assert False, "not supposed to get here" @@ -200,16 +160,16 @@ def generate_unroll_or_ilp_code(cgs, kernel, sched_index, exec_domain): # {{{ parallel loop -def generate_parallel_loop_dim_code(cgs, kernel, sched_index, exec_domain): +def generate_parallel_loop_dim_code(kernel, sched_index, exec_domain): from loopy.isl import make_slab - ccm = cgs.c_code_mapper + ccm = exec_domain.c_code_mapper space = kernel.space iname = kernel.schedule[sched_index].iname tag = kernel.iname_to_tag.get(iname) lb_cns_orig, ub_cns_orig, slabs = get_slab_decomposition( - cgs, kernel, sched_index, exec_domain) + kernel, sched_index, exec_domain) # For a parallel loop dimension, the global loop bounds are # automatically obeyed--simply because no work items are launched @@ -242,8 +202,7 @@ def generate_parallel_loop_dim_code(cgs, kernel, sched_index, exec_domain): domain=kernel.domain.intersect(slab)) result.append( add_comment(cmt, - build_loop_nest(cgs, new_kernel, sched_index+1, - exec_domain))) + build_loop_nest(new_kernel, sched_index+1, exec_domain))) from loopy.codegen import gen_code_block return gen_code_block(result, is_alternatives=True) @@ -252,15 +211,15 @@ def generate_parallel_loop_dim_code(cgs, kernel, sched_index, exec_domain): # {{{ sequential loop -def generate_sequential_loop_dim_code(cgs, kernel, sched_index, exec_domain): +def generate_sequential_loop_dim_code(kernel, sched_index, exec_domain): - ccm = cgs.c_code_mapper + ccm = exec_domain.c_code_mapper space = kernel.space iname = kernel.schedule[sched_index].iname tag = kernel.iname_to_tag.get(iname) lb_cns_orig, ub_cns_orig, slabs = get_slab_decomposition( - cgs, kernel, sched_index, exec_domain) + kernel, sched_index, exec_domain) result = [] nums_of_conditionals = [] @@ -271,7 +230,7 @@ def generate_sequential_loop_dim_code(cgs, kernel, sched_index, exec_domain): cmt = None new_exec_domain = exec_domain.intersect(slab) - inner = build_loop_nest(cgs, kernel, sched_index+1, + inner = build_loop_nest(kernel, sched_index+1, new_exec_domain) from loopy.codegen.bounds import wrap_in_for_from_constraints diff --git a/loopy/codegen/prefetch.py b/loopy/codegen/prefetch.py index 22af66d33..155927e4b 100644 --- a/loopy/codegen/prefetch.py +++ b/loopy/codegen/prefetch.py @@ -105,24 +105,21 @@ def make_fetch_loop_nest(flnd, fetch_dim_idx, pf_dim_exprs, iname_subst_map, # done, return from pymbolic.primitives import Variable, Subscript - from pymbolic.mapper.stringifier import PREC_NONE result = Assign( pf.name + "".join("[%s]" % ccm(dexpr) for dexpr in pf_dim_exprs), no_pf_ccm( Subscript( Variable(pf.input_vector), - substitute(pf.index_expr, iname_subst_map)), - PREC_NONE)) - - def my_ccm(expr): - return ccm(substitute(expr, iname_subst_map)) + substitute(pf.index_expr, iname_subst_map)))) from pymbolic.mapper.dependency import DependencyMapper check_vars = [v.name for v in DependencyMapper()(pf.index_expr)] from loopy.codegen.bounds import wrap_in_bounds_checks - return wrap_in_bounds_checks(my_ccm, pf.kernel.domain, + return wrap_in_bounds_checks( + ccm.copy_and_assign_many(iname_subst_map), + pf.kernel.domain, check_vars, implemented_domain, result) fetch_inames = pf.fetch_dims[fetch_dim_idx] @@ -280,12 +277,12 @@ def make_fetch_loop_nest(flnd, fetch_dim_idx, pf_dim_exprs, iname_subst_map, # }}} -def generate_prefetch_code(cgs, kernel, sched_index, exec_domain): +def generate_prefetch_code(kernel, sched_index, exec_domain): implemented_domain = exec_domain.implemented_domain from cgen import Statement as S, Line, Comment - ccm = cgs.c_code_mapper + ccm = exec_domain.c_code_mapper scheduled_pf = kernel.schedule[sched_index] pf = kernel.prefetch[ @@ -482,7 +479,7 @@ def generate_prefetch_code(cgs, kernel, sched_index, exec_domain): from loopy.codegen.dispatch import build_loop_nest new_block.extend([Line(), - build_loop_nest(cgs, kernel, sched_index+1, exec_domain)]) + build_loop_nest(kernel, sched_index+1, exec_domain)]) return gen_code_block(new_block) diff --git a/loopy/compiled.py b/loopy/compiled.py index bd72ecdaf..a2b4210d6 100644 --- a/loopy/compiled.py +++ b/loopy/compiled.py @@ -44,7 +44,7 @@ class CompiledKernel: from pymbolic import compile if size_args is None: - self.size_args = [arg.name for arg in kernel.args if isinstance(arg, ScalarArg)] + self.size_args = kernel.scalar_loop_args else: self.size_args = size_args diff --git a/loopy/kernel.py b/loopy/kernel.py index f1c863838..30aa1c515 100644 --- a/loopy/kernel.py +++ b/loopy/kernel.py @@ -6,6 +6,8 @@ import numpy as np from pytools import Record, memoize_method from pymbolic.mapper.dependency import DependencyMapper import pyopencl as cl +import islpy as isl +from islpy import dim_type @@ -32,9 +34,17 @@ class ArrayArg: raise ValueError("can only specify one of shape and strides") if strides is not None: + if isinstance(strides, str): + from pymbolic import parse + strides = parse(strides) + strides = tuple(strides) if shape is not None: + if isinstance(shape, str): + from pymbolic import parse + shape = parse(shape) + from pyopencl.compyte.array import ( f_contiguous_strides, c_contiguous_strides) @@ -67,7 +77,7 @@ class ImageArg: class ScalarArg: - def __init__(self, name, dtype, approximately): + def __init__(self, name, dtype, approximately=None): self.name = name self.dtype = np.dtype(dtype) self.approximately = approximately @@ -176,34 +186,57 @@ class LoopKernel(Record): :ivar register_prefetch: :ivar name: :ivar preamble: + :ivar assumptions: the initial implemented_domain, captures assumptions + on the parameters. (an isl.Set) + :ivar iname_slab_increments: a dictionary mapping inames to (lower_incr, + upper_incr) tuples that will be separated out in the execution to generate + 'bulk' slabs with fewer conditionals. """ def __init__(self, device, domain, instructions, args=None, prefetch={}, schedule=None, register_prefetch=None, name="loopy_kernel", iname_to_tag={}, is_divisible=lambda dividend, divisor: False, - preamble=None): + preamble=None, assumptions=None, + iname_slab_increments={}): """ :arg domain: a :class:`islpy.BasicSet`, or a string parseable to a basic set by the isl. Example: "{[i,j]: 0<=i < 10 and 0<= j < 9}" :arg iname_to_tag: a map from loop domain variables to subclasses of :class:`IndexTag` """ - from pymbolic import parse - def parse_if_necessary(v): - if isinstance(v, str): - return parse(v) - else: - return v + def parse_if_necessary(insn): + from pymbolic import parse + if isinstance(insn, str): + lhs, rhs = insn.split("=") + elif isinstance(insn, tuple): + lhs, rhs = insn + + if isinstance(lhs, str): + lhs = parse(lhs) + if isinstance(rhs, str): + from loopy.symbolic import CSERealizer + rhs = CSERealizer()(parse(rhs)) + + return lhs, rhs if isinstance(domain, str): - import islpy as isl ctx = isl.Context() domain = isl.Set.read_from_str(ctx, domain, nparam=-1) - insns = [ - (parse_if_necessary(lvalue), - parse_if_necessary(expr)) - for lvalue, expr in instructions] + insns = [parse_if_necessary(insn) for insn in instructions] + + if assumptions is None: + assumptions = isl.Set.universe(domain.get_dim()) + elif isinstance(assumptions, str): + s = domain.get_dim() + assumptions = isl.BasicSet.read_from_str(domain.get_ctx(), + "[%s] -> {[%s]: %s}" + % (",".join(s.get_name(dim_type.param, i) + for i in range(s.size(dim_type.param))), + ",".join(s.get_name(dim_type.set, i) + for i in range(s.size(dim_type.set))), + assumptions), + nparam=-1) Record.__init__(self, device=device, args=args, domain=domain, instructions=insns, @@ -213,7 +246,9 @@ class LoopKernel(Record): (iname, parse_tag(tag)) for iname, tag in iname_to_tag.iteritems()), is_divisible=is_divisible, - preamble=preamble) + preamble=preamble, + assumptions=assumptions, + iname_slab_increments=iname_slab_increments) # FIXME: assert empty intersection of loop vars and args # FIXME: assert that isl knows about loop parameters @@ -236,12 +271,16 @@ class LoopKernel(Record): def arg_dict(self): return dict((arg.name, arg) for arg in self.args) + @property @memoize_method - def scalar_args(self): + def scalar_loop_args(self): if self.args is None: - return set() + return [] else: - return set(arg.name for arg in self.args if isinstance(arg, ScalarArg)) + loop_arg_names = [self.space.get_name(dim_type.param, i) + for i in range(self.space.size(dim_type.param))] + return [arg.name for arg in self.args if isinstance(arg, ScalarArg) + if arg.name in loop_arg_names] @memoize_method def all_inames(self): @@ -332,7 +371,7 @@ class LoopKernel(Record): input_vectors.update( set(v.name for v in dm(expr))) - return input_vectors - self.all_inames() - self.scalar_args() + return input_vectors - self.all_inames() - set(self.scalar_loop_args) @memoize_method def output_vectors(self): @@ -343,7 +382,7 @@ class LoopKernel(Record): output_vectors.update( set(v.name for v in dm(lvalue))) - return output_vectors - self.all_inames() - self.scalar_args() + return output_vectors - self.all_inames() - self.scalar_loop_args def _subst_insns(self, old_var, new_expr): from pymbolic.mapper.substitutor import substitute @@ -393,7 +432,8 @@ class LoopKernel(Record): def split_dimension(self, name, inner_length, padded_length=None, outer_name=None, inner_name=None, - outer_tag=None, inner_tag=None): + outer_tag=None, inner_tag=None, + outer_slab_increments=(0, -1)): outer_tag = parse_tag(outer_tag) inner_tag = parse_tag(inner_tag) @@ -434,33 +474,41 @@ class LoopKernel(Record): from islpy import dim_type outer_var_nr = self.space.size(dim_type.set) inner_var_nr = self.space.size(dim_type.set)+1 - new_domain = self.domain.add_dims(dim_type.set, 2) - new_domain.set_dim_name(dim_type.set, outer_var_nr, outer_name) - new_domain.set_dim_name(dim_type.set, inner_var_nr, inner_name) - import islpy as isl - from loopy.isl import make_slab + def process_set(s): + s = s.add_dims(dim_type.set, 2) + s.set_dim_name(dim_type.set, outer_var_nr, outer_name) + s.set_dim_name(dim_type.set, inner_var_nr, inner_name) + + from loopy.isl import make_slab + + space = s.get_dim() + inner_constraint_set = ( + make_slab(space, inner_name, 0, inner_length) + # name = inner + length*outer + .add_constraint(isl.Constraint.eq_from_names( + space, {name:1, inner_name: -1, outer_name:-inner_length}))) - space = new_domain.get_dim() - inner_constraint_set = ( - make_slab(space, inner_name, 0, inner_length) - # name = inner + length*outer - .add_constraint(isl.Constraint.eq_from_names( - space, {name:1, inner_name: -1, outer_name:-inner_length}))) + name_dim_type, name_idx = space.get_var_dict()[name] + return (s + .intersect(inner_constraint_set) + .project_out(name_dim_type, name_idx, 1)) - name_dim_type, name_idx = space.get_var_dict()[name] - new_domain = (new_domain - .intersect(inner_constraint_set) - .project_out(name_dim_type, name_idx, 1)) + new_domain = process_set(self.domain) + new_assumptions = process_set(self.assumptions) from pymbolic import var inner = var(inner_name) outer = var(outer_name) new_loop_index = inner + outer*inner_length + iname_slab_increments = self.iname_slab_increments.copy() + iname_slab_increments[outer_name] = outer_slab_increments return (self .substitute(name, new_loop_index) - .copy(domain=new_domain, iname_to_tag=new_iname_to_tag)) + .copy(domain=new_domain, iname_to_tag=new_iname_to_tag, + assumptions=new_assumptions, + iname_slab_increments=iname_slab_increments)) def get_problems(self, parameters, emit_warnings=True): """ diff --git a/loopy/prefetch.py b/loopy/prefetch.py index 49330236c..98aeed948 100644 --- a/loopy/prefetch.py +++ b/loopy/prefetch.py @@ -9,7 +9,7 @@ from islpy import dim_type # {{{ register prefetches class RegisterPrefetch(Record): - __slots__ = ["subscript_expr", "new_name"] + __slots__ = ["subexprs", "new_names"] def insert_register_prefetches(kernel): reg_pf = {} @@ -54,7 +54,7 @@ def insert_register_prefetches(kernel): reg_pf[iexpr] = new_name schedule.insert(sched_index, RegisterPrefetch( - index_expr=iexpr, new_name=new_name)) + subexprs=[iexpr], new_names=[new_name])) sched_index += 1 else: i += 1 @@ -178,7 +178,7 @@ class LocalMemoryPrefetch(Record): from pymbolic.mapper.dependency import DependencyMapper return set(var.name for var in DependencyMapper()(self.index_expr) - ) - set(self.all_inames()) - self.kernel.scalar_args() + ) - set(self.all_inames()) - set(self.kernel.scalar_loop_args) def hash(self): return (hash(type(self)) ^ hash(self.input_vector) diff --git a/loopy/symbolic.py b/loopy/symbolic.py index 2e2587f1d..b928d4b31 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -2,7 +2,7 @@ from __future__ import division -from pymbolic.mapper import CombineMapper, RecursiveMapper +from pymbolic.mapper import CombineMapper, RecursiveMapper, IdentityMapper from pymbolic.mapper.c_code import CCodeMapper from pymbolic.mapper.stringifier import PREC_NONE import numpy as np @@ -120,7 +120,8 @@ class VariableIndexExpressionCollector(CombineMapper): # {{{ C code mapper class LoopyCCodeMapper(CCodeMapper): - def __init__(self, kernel, no_prefetch=False): + def __init__(self, kernel, no_prefetch=False, cse_name_list=[], + var_subst_map={}): def constant_mapper(c): if isinstance(c, float): # FIXME: type-variable @@ -128,11 +129,39 @@ class LoopyCCodeMapper(CCodeMapper): else: return repr(c) - CCodeMapper.__init__(self, constant_mapper=constant_mapper) + CCodeMapper.__init__(self, constant_mapper=constant_mapper, + cse_name_list=cse_name_list) self.kernel = kernel + self.var_subst_map = var_subst_map.copy() + self.no_prefetch = no_prefetch + def copy(self, var_subst_map=None, cse_name_list=None): + if var_subst_map is None: + var_subst_map = self.var_subst_map + if cse_name_list is None: + cse_name_list = self.cse_name_list + return LoopyCCodeMapper(self.kernel, no_prefetch=self.no_prefetch, + cse_name_list=cse_name_list, var_subst_map=var_subst_map) + + def copy_and_assign(self, name, value): + var_subst_map = self.var_subst_map.copy() + var_subst_map[name] = value + return self.copy(var_subst_map=var_subst_map) + + def copy_and_assign_many(self, assignments): + var_subst_map = self.var_subst_map.copy() + var_subst_map.update(assignments) + return self.copy(var_subst_map=var_subst_map) + + def map_variable(self, expr, prec): + if expr.name in self.var_subst_map: + return " /* %s */ %s" % ( + expr.name, self.rec(self.var_subst_map[expr.name], prec)) + else: + return CCodeMapper.map_variable(self, expr, prec) + def map_subscript(self, expr, enclosing_prec): from pymbolic.primitives import Variable if (not self.no_prefetch @@ -143,11 +172,11 @@ class LoopyCCodeMapper(CCodeMapper): except KeyError: pass else: - from pymbolic.mapper.stringifier import PREC_SUM + from pymbolic import var return pf.name+"".join( - "[%s - %s]" % (iname, self.rec( - pf.dim_bounds_by_iname[iname][0], - PREC_SUM)) + "[%s]" % self.rec( + var(iname) - pf.dim_bounds_by_iname[iname][0], + PREC_NONE) for iname in pf.all_inames()) if isinstance(expr.aggregate, Variable): @@ -252,6 +281,26 @@ def constraint_to_expr(cns, except_name=None): # }}} +# {{{ CSE "realizer" + +class CSERealizer(IdentityMapper): + """Looks for invocations of a function called 'cse' and turns those into + CommonSubexpression objects. + """ + + def map_call(self, expr): + from pymbolic import Variable + if isinstance(expr.function, Variable) and expr.function.name == "cse": + from pymbolic.primitives import CommonSubexpression + if len(expr.parameters) == 1: + return CommonSubexpression(expr.parameters[0]) + else: + raise TypeError("cse takes exactly one argument") + else: + return IdentityMapper.map_call(self, expr) + +# }}} + diff --git a/test/test_matmul.py b/test/test_matmul.py index 7aa72e646..c512640c4 100644 --- a/test/test_matmul.py +++ b/test/test_matmul.py @@ -2,6 +2,7 @@ import numpy as np import numpy.linalg as la import pyopencl as cl import pyopencl.array as cl_array +import pyopencl.clrandom as cl_random import loopy as lp from pyopencl.tools import pytest_generate_tests_for_pyopencl \ @@ -47,7 +48,12 @@ def check_error(refsol, sol): if not DO_CHECK: return - rel_err = la.norm(refsol-sol, "fro")/la.norm(refsol, "fro") + if sol.shape == 2: + norm_order = "fro" + else: + norm_order = 2 + + rel_err = la.norm(refsol-sol, norm_order)/la.norm(refsol, norm_order) if rel_err > 1e-5 or np.isinf(rel_err) or np.isnan(rel_err): if 1: import matplotlib.pyplot as pt @@ -86,22 +92,20 @@ def test_axpy(ctx_factory): queue = cl.CommandQueue(ctx, properties=cl.command_queue_properties.PROFILING_ENABLE) - n = get_suitable_size(ctx) - from pymbolic import var - x, y, z, n_sym, i = [var(s) for s in "xyzni"] - - n_approx = 10**6 + n = get_suitable_size(ctx)**3 knl = lp.LoopKernel(ctx.devices[0], "[n] -> {[i]: 0<=i<n}", [ - (z[i], x[i]+y[i]) # FIXME: Add scalars + "z[i] = a*x[i]+b*y[i]" ], [ - lp.ArrayArg("x", dtype, shape=(n,)), - lp.ArrayArg("y", dtype, shape=(n,)), - lp.ArrayArg("z", dtype, shape=(n,)), - lp.ScalarArg("n", np.int32, approximately=n_approx), + lp.ScalarArg("a", dtype), + lp.ArrayArg("x", dtype, shape="n,"), + lp.ScalarArg("b", dtype), + lp.ArrayArg("y", dtype, shape="n,"), + lp.ArrayArg("z", dtype, shape="n,"), + lp.ScalarArg("n", np.int32, approximately=n), ], name="matmul") @@ -109,27 +113,26 @@ def test_axpy(ctx_factory): block_size = 256 knl = lp.split_dimension(knl, "i", unroll*block_size, outer_tag="g.0") knl = lp.split_dimension(knl, "i_inner", block_size, outer_tag="unr", inner_tag="l.0") - assert knl.get_problems({"n": n_approx})[0] <= 2 + assert knl.get_problems({"n": n})[0] <= 2 kernel_gen = (lp.insert_register_prefetches(knl) for knl in lp.generate_loop_schedules(knl)) - #a = make_well_conditioned_dev_matrix(queue, n, dtype=dtype, order=order) - #b = make_well_conditioned_dev_matrix(queue, n, dtype=dtype, order=order) - #c = cl_array.empty_like(a) - #refsol = np.dot(a.get(), b.get()) + a = cl_random.rand(queue, n, dtype=dtype) + b = cl_random.rand(queue, n, dtype=dtype) + c = cl_array.empty_like(a) + refsol = (2*a+3*b).get() def launcher(kernel, gsize, lsize, check): - #evt = kernel(queue, gsize(), lsize(), a.data, b.data, c.data, - #g_times_l=True) + evt = kernel(queue, gsize(n), lsize(n), 2, a.data, 3, b.data, c.data, n, + g_times_l=True) - #if check: - #check_error(refsol, c.get()) + if check: + check_error(refsol, c.get()) - #return evt - pass + return evt - lp.drive_timing_run(kernel_gen, queue, launcher, 2*n**3) + lp.drive_timing_run(kernel_gen, queue, launcher, 5*n) @@ -143,13 +146,11 @@ def test_plain_matrix_mul(ctx_factory): properties=cl.command_queue_properties.PROFILING_ENABLE) n = get_suitable_size(ctx) - from pymbolic import var - a, b, c, i, j, k, n_sym = [var(s) for s in "abcijkn"] knl = lp.LoopKernel(ctx.devices[0], "{[i,j,k]: 0<=i,j,k<%d}" % n, [ - (c[i, j], a[i, k]*b[k, j]) + "c[i, j] = a[i, k]*b[k, j]" ], [ lp.ArrayArg("a", dtype, shape=(n, n), order=order), @@ -196,13 +197,11 @@ def test_image_matrix_mul(ctx_factory): properties=cl.command_queue_properties.PROFILING_ENABLE) n = get_suitable_size(ctx) - from pymbolic import var - a, b, c, i, j, k, n_sym = [var(s) for s in "abcijkn"] knl = lp.LoopKernel(ctx.devices[0], "{[i,j,k]: 0<=i,j,k<%d}" % n, [ - (c[i, j], a[i, k]*b[k, j]) + "c[i, j] = a[i, k]*b[k, j]" ], [ lp.ImageArg("a", dtype, 2), @@ -251,13 +250,11 @@ def test_image_matrix_mul_ilp(ctx_factory): properties=cl.command_queue_properties.PROFILING_ENABLE) n = get_suitable_size(ctx) - from pymbolic import var - a, b, c, i, j, k, n_sym = [var(s) for s in "abcijkn"] knl = lp.LoopKernel(ctx.devices[0], "{[i,j,k]: 0<=i,j,k<%d}" % n, [ - (c[i, j], a[i, k]*b[k, j]) + "c[i, j] = a[i, k]*b[k, j]" ], [ lp.ImageArg("a", dtype, 2), @@ -312,18 +309,16 @@ def test_fancy_matrix_mul(ctx_factory): order = "C" n = get_suitable_size(ctx) - from pymbolic import var - a, b, c, i, j, k, n_sym = [var(s) for s in "abcijkn"] knl = lp.LoopKernel(ctx.devices[0], "[n] -> {[i,j,k]: 0<=i,j,k<n }", [ - (c[i, j], a[i, k]*b[k, j]) + "c[i, j] = a[i, k]*b[k, j]" ], [ - lp.ArrayArg("a", dtype, shape=(n_sym, n_sym), order=order), - lp.ArrayArg("b", dtype, shape=(n_sym, n_sym), order=order), - lp.ArrayArg("c", dtype, shape=(n_sym, n_sym), order=order), + lp.ArrayArg("a", dtype, shape="(n, n)", order=order), + lp.ArrayArg("b", dtype, shape="(n, n)", order=order), + lp.ArrayArg("c", dtype, shape="(n, n)", order=order), lp.ScalarArg("n", np.int32, approximately=1000), ], name="fancy_matmul") -- GitLab