Scalar min/max?

[nicknytko]

Hello, I'm back with hopefully another obvious question 🙂

I'm trying to take the minimum of two scalars in a kernel using min(), but I am running into a parse error on the expression. Are the semantics of min and max not the same as in Python? Minimal example to illustrate what I'm trying to do:

import loopy

domain = "{ [i]: 0 <= i < n }"
body = """
out[i] = min(a[i], b[i])
"""

kernel = loopy.make_kernel(domain, body, lang_version=(2018, 2))
print(kernel)

The stacktrace that I'm getting:

While parsing right hand side 'min(a[i], b[i])', the following error occurred:
Traceback (most recent call last):
  File "/Users/nnytko2/Desktop/loopy_min.py", line 10, in <module>
    kernel = loopy.make_kernel(domain, body, lang_version=(2018, 2))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/loopy/kernel/creation.py", line 2569, in make_kernel
    tunit = make_function(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/loopy/kernel/creation.py", line 2394, in make_function
    parse_instructions(instructions, defines)
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/loopy/kernel/creation.py", line 1000, in parse_instructions
    insn, insn_inames_to_dup = parse_insn(
                               ^^^^^^^^^^^
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/loopy/kernel/creation.py", line 505, in parse_insn
    rhs = parse(groups["rhs"])
          ^^^^^^^^^^^^^^^^^^^^
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/loopy/symbolic.py", line 1715, in parse
    FunctionToPrimitiveMapper()(LoopyParser()(expr_str)))
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/pymbolic/mapper/__init__.py", line 141, in __call__
    result = method(expr, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/loopy/symbolic.py", line 1619, in map_call
    return self._parse_reduction(operation, inames, red_exprs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/nnytko2/.pyenv/versions/3.11.2/lib/python3.11/site-packages/loopy/symbolic.py", line 1546, in _parse_reduction
    raise TypeError("iname argument to reduce() must be a symbol "
TypeError: iname argument to reduce() must be a symbol or a list/tuple of symbols

[inducer]

Similar to numpy, min is a reduction, so you could say min(i, a[i]) to compute the minimum of an array. In numpy, you would use minimum for the minimum of two numbers. In loopy, I don't think we can parse minimum just yet, so the way to create the corresponding (pymbolic) expression node is kind of clunky ATM:

import pyopencl as cl

import numpy as np
import loopy as lp
import pymbolic.primitives as p
from pymbolic import parse

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

domain = "{ [i]: 0 <= i < n }"
body = [
    lp.Assignment("out[i]", p.Min((parse("a[i]"), parse("b[i]"))))
    ]

kernel = lp.make_kernel(domain, body, lang_version=(2018, 2))

a = np.random.randn(1000)
b = np.random.randn(1000)
_evt, (res,) = kernel(queue, a=a, b=b)
assert np.array_equal(res, np.minimum(a, b))

Patches welcome. The code to do this better would go here:

loopy/loopy/symbolic.py

Lines 1563 to 1622 in 524f2c4

	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)

[nicknytko]

Fixed with patch in #795