def map_call(self, expr):
from loopy.library.reduction import parse_reduction_op
if not isinstance(expr.function, p.Variable):
return IdentityMapper.map_call(self, expr)
name = expr.function.name
if name == "cse":
if len(expr.parameters) in [1, 2]:
if len(expr.parameters) == 2:
if not isinstance(expr.parameters[1], p.Variable):
raise TypeError("second argument to cse() must be a symbol")
tag = expr.parameters[1].name
else:
tag = None
return p.CommonSubexpression(
self.rec(expr.parameters[0]), tag)
else:
raise TypeError("cse takes two arguments")
elif name in ["reduce", "simul_reduce"]:
if len(expr.parameters) >= 3:
operation, inames = expr.parameters[:2]
red_exprs = expr.parameters[2:]
operation = parse_reduction_op(str(operation))
return self._parse_reduction(operation, inames,
tuple(self.rec(red_expr) for red_expr in red_exprs),
allow_simultaneous=(name == "simul_reduce"))
else:
raise TypeError("invalid 'reduce' calling sequence")
elif name == "if":
if len(expr.parameters) == 3:
from pymbolic.primitives import If
return If(*tuple(self.rec(p) for p in expr.parameters))
else:
raise TypeError("if takes three arguments")
else:
# see if 'name' is an existing reduction op
operation = parse_reduction_op(name)
if operation:
# arg_count counts arguments but not inames
if len(expr.parameters) != 1 + operation.arg_count:
raise RuntimeError("invalid invocation of "
"reduction operation '%s': expected %d arguments, "
"got %d instead" % (expr.function.name,
1 + operation.arg_count,
len(expr.parameters)))
inames = expr.parameters[0]
red_exprs = tuple(self.rec(param) for param in expr.parameters[1:])
return self._parse_reduction(operation, inames, red_exprs)
else:
return IdentityMapper.map_call(self, expr)
def map_call(self, expr):
from loopy.library.reduction import parse_reduction_op
from pymbolic.primitives import Variable
if not isinstance(expr.function, Variable):
return IdentityMapper.map_call(self, expr)
name = expr.function.name
if name == "cse":
from pymbolic.primitives import CommonSubexpression
if len(expr.parameters) in [1, 2]:
if len(expr.parameters) == 2:
if not isinstance(expr.parameters[1], Variable):
raise TypeError("second argument to cse() must be a symbol")
tag = expr.parameters[1].name
else:
tag = None
return CommonSubexpression(
self.rec(expr.parameters[0]), tag)
else:
raise TypeError("cse takes two arguments")
elif name in ["reduce", "simul_reduce"]:
if len(expr.parameters) == 3:
operation, inames, red_expr = expr.parameters
if not isinstance(operation, Variable):
raise TypeError("operation argument to reduce() "
"must be a symbol")
operation = parse_reduction_op(operation.name)
return self._parse_reduction(operation, inames, self.rec(red_expr),
allow_simultaneous=(name == "simul_reduce"))
else:
raise TypeError("invalid 'reduce' calling sequence")
elif name == "if":
if len(expr.parameters) == 3:
from pymbolic.primitives import If
return If(*expr.parameters)
else:
raise TypeError("if takes three arguments")
else:
# see if 'name' is an existing reduction op
operation = parse_reduction_op(name)
if operation:
if len(expr.parameters) != 2:
raise RuntimeError("invalid invocation of "
"reduction operation '%s'" % expr.function.name)
inames, red_expr = expr.parameters
return self._parse_reduction(operation, inames, self.rec(red_expr))
else:
return IdentityMapper.map_call(self, expr)
def map_call(self, expr):
from loopy.library.reduction import parse_reduction_op
from pymbolic.primitives import Variable
if not isinstance(expr.function, Variable):
return IdentityMapper.map_call(self, expr)
name = expr.function.name
if name == "cse":
from pymbolic.primitives import CommonSubexpression
if len(expr.parameters) in [1, 2]:
if len(expr.parameters) == 2:
if not isinstance(expr.parameters[1], Variable):
raise TypeError("second argument to cse() must be a symbol")
tag = expr.parameters[1].name
else:
tag = None
return CommonSubexpression(
self.rec(expr.parameters[0]), tag)
else:
raise TypeError("cse takes two arguments")
elif name in ["reduce", "simul_reduce"]:
if len(expr.parameters) == 3:
operation, inames, red_expr = expr.parameters
if not isinstance(operation, Variable):
raise TypeError("operation argument to reduce() "
"must be a symbol")
operation = parse_reduction_op(operation.name)
return self._parse_reduction(operation, inames, self.rec(red_expr),
allow_simultaneous=(name == "simul_reduce"))
else:
raise TypeError("invalid 'reduce' calling sequence")
elif name == "if":
if len(expr.parameters) == 3:
from pymbolic.primitives import If
return If(*expr.parameters)
else:
raise TypeError("if takes three arguments")
else:
# see if 'name' is an existing reduction op
operation = parse_reduction_op(name)
if operation:
if len(expr.parameters) != 2:
raise RuntimeError("invalid invocation of "
"reduction operation '%s'" % expr.function.name)
inames, red_expr = expr.parameters
return self._parse_reduction(operation, inames, self.rec(red_expr))
else:
return IdentityMapper.map_call(self, expr)
def __init__(self, operation, inames, expr, allow_simultaneous=False):
if isinstance(inames, str):
inames = tuple(iname.strip() for iname in inames.split(","))
elif isinstance(inames, Variable):
inames = (inames,)
assert isinstance(inames, tuple)
def strip_var(iname):
if isinstance(iname, Variable):
iname = iname.name
assert isinstance(iname, str)
return iname
inames = tuple(strip_var(iname) for iname in inames)
if isinstance(operation, str):
from loopy.library.reduction import parse_reduction_op
operation = parse_reduction_op(operation)
from loopy.library.reduction import ReductionOperation
assert isinstance(operation, ReductionOperation)
self.operation = operation
self.inames = inames
self.expr = expr
self.allow_simultaneous = allow_simultaneous
def __init__(self, operation, inames, expr, allow_simultaneous=False):
if isinstance(inames, str):
inames = tuple(iname.strip() for iname in inames.split(","))
elif isinstance(inames, Variable):
inames = (inames,)
assert isinstance(inames, tuple)
def strip_var(iname):
if isinstance(iname, Variable):
iname = iname.name
assert isinstance(iname, str)
return iname
inames = tuple(strip_var(iname) for iname in inames)
if isinstance(operation, str):
from loopy.library.reduction import parse_reduction_op
operation = parse_reduction_op(operation)
from loopy.library.reduction import ReductionOperation
assert isinstance(operation, ReductionOperation)
self.operation = operation
self.inames = inames
self.expr = expr
self.allow_simultaneous = allow_simultaneous
def map_matrix_product(self, expr: MatrixProduct,
state: CodeGenState) -> ImplementedResult:
if expr in state.results:
return state.results[expr]
x1_result = self.rec(expr.x1, state)
x2_result = self.rec(expr.x2, state)
loopy_expr_context = LoopyExpressionContext(state,
num_indices=expr.ndim)
loopy_expr_context.reduction_bounds["_r0"] = (0, expr.x2.shape[0])
# Figure out inames.
x1_inames = []
for i in range(expr.x1.ndim):
if i == expr.x1.ndim - 1:
x1_inames.append(var("_r0"))
else:
x1_inames.append(var(f"_{i}"))
x2_inames = []
for i in range(expr.x2.ndim):
if i == 0:
x2_inames.append(var("_r0"))
else:
offset = i + len(x1_inames) - 2
x2_inames.append(var(f"_{offset}"))
inner_expr = x1_result.to_loopy_expression(tuple(x1_inames),
loopy_expr_context)
inner_expr *= x2_result.to_loopy_expression(tuple(x2_inames),
loopy_expr_context)
import loopy.library.reduction as red
loopy_expr = lp.Reduction(operation=red.parse_reduction_op("sum"),
inames=("_r0", ),
expr=inner_expr,
allow_simultaneous=False)
inlined_result = InlinedResult.from_loopy_expression(
loopy_expr, loopy_expr_context)
output_name = state.var_name_gen("matmul")
insn_id = add_store(output_name,
expr,
inlined_result,
state,
output_to_temporary=True)
result = StoredResult(output_name, expr.ndim, frozenset([insn_id]))
state.results[expr] = result
return result
def __init__(self, operation, inames, expr, allow_simultaneous=False):
if isinstance(inames, str):
inames = tuple(iname.strip() for iname in inames.split(","))
elif isinstance(inames, p.Variable):
inames = (inames, )
assert isinstance(inames, tuple)
def strip_var(iname):
if isinstance(iname, p.Variable):
iname = iname.name
assert isinstance(iname, str)
return iname
inames = tuple(strip_var(iname) for iname in inames)
if isinstance(operation, str):
from loopy.library.reduction import parse_reduction_op
operation = parse_reduction_op(operation)
from loopy.library.reduction import ReductionOperation
assert isinstance(operation, ReductionOperation)
from loopy.diagnostic import LoopyError
if operation.arg_count > 1:
from pymbolic.primitives import Call
if not isinstance(expr, (tuple, Reduction, Call)):
raise LoopyError("reduction argument must be one of "
"a tuple, reduction, or call; "
"got '%s'" % type(expr).__name__)
else:
# Sanity checks
if isinstance(expr, tuple):
raise LoopyError("got a tuple argument to a scalar reduction")
elif isinstance(expr, Reduction) and expr.is_tuple_typed:
raise LoopyError(
"got a tuple typed argument to a scalar reduction")
self.operation = operation
self.inames = inames
self.expr = expr
self.allow_simultaneous = allow_simultaneous