Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions arraycontext/impl/pytato/compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,8 +218,10 @@ def _get_f_placeholder_args(arg, kw, arg_id_to_name, actx):
:attr:`BaseLazilyCompilingFunctionCaller.f`.
"""
if np.isscalar(arg):
from pytato.tags import ForceValueArgTag
name = arg_id_to_name[kw,]
return pt.make_placeholder(name, (), np.dtype(type(arg)))
return pt.make_placeholder(name, (), np.dtype(type(arg)),
tags=frozenset({ForceValueArgTag()}))
elif isinstance(arg, pt.Array):
name = arg_id_to_name[kw,]
# Transform the DAG to give metadata inference a chance to do its job
Expand Down Expand Up @@ -533,9 +535,8 @@ def _args_to_device_buffers(actx, input_id_to_name_in_program, arg_id_to_arg):
for arg_id, arg in arg_id_to_arg.items():
if np.isscalar(arg):
if isinstance(actx, PytatoPyOpenCLArrayContext):
import pyopencl.array as cla
arg = cla.to_device(actx.queue, np.array(arg),
allocator=actx.allocator)
# Scalar kernel args are passed as lp.ValueArgs
pass
elif isinstance(actx, PytatoJAXArrayContext):
import jax
arg = jax.device_put(arg)
Expand Down
27 changes: 27 additions & 0 deletions test/test_pytato_arraycontext.py
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,33 @@ def test_transfer(actx_factory):
# }}}


def test_pass_args_compiled_func(actx_factory):
import numpy as np

import loopy as lp
import pyopencl as cl
import pyopencl.array
import pytato as pt

def twice(x, y, a):
return 2 * x * y * a

actx = _PytatoPyOpenCLArrayContextForTests(actx_factory().queue)

dev_scalar = pt.make_data_wrapper(cl.array.to_device(actx.queue, np.float64(23)))

f = actx.compile(twice)

assert actx.to_numpy(f(99.0, np.float64(2.0), dev_scalar)) == 2*23*99*2

compiled_func, = f.program_cache.values()
ep = compiled_func.pytato_program.program.t_unit.default_entrypoint

assert isinstance(ep.arg_dict["_actx_in_0"], lp.ValueArg)
assert isinstance(ep.arg_dict["_actx_in_1"], lp.ValueArg)
assert isinstance(ep.arg_dict["_actx_in_2"], lp.ArrayArg)


if __name__ == "__main__":
import sys
if len(sys.argv) > 1:
Expand Down
Loading