match_args_usage_pattern.py

import loopy as lp
import pyopencl as cl
import pyopencl.clrandom as cl_random
import numpy as np
import numpy.linalg as la

from loopy.transform.register_callable import _match_caller_callee_argument_dimension

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

x1 = cl_random.rand(queue, (3, 2), dtype=np.float64)
x2 = cl_random.rand(queue, (6, ), dtype=np.float64)

callee1 = lp.make_kernel(
        "{[i]: 0<=i<6}",
        """
        a[i] = 2*abs(b[i])
        """)

callee2 = lp.make_kernel(
        "{[i, j]: 0<=i<3 and 0 <= j < 2}",
        """
        a[i, j] = 3*b[i, j]
        """)

knl = lp.make_kernel(
        "{[i, j, k]:  0<= i < 3 and 0 <= j < 2 and 0 <= k < 6}",
        """
        [i, j]: y1[i, j] = callee_fn1([i, j]: x1[i, j])
        [k]: y2[k] = callee_fn2([k]: x2[k])
        """)

knl = lp.register_callable_kernel(knl, 'callee_fn1', callee1)
knl = lp.register_callable_kernel(knl, 'callee_fn2', callee2)
knl = _match_caller_callee_argument_dimension(knl, 'callee_fn1')
knl = _match_caller_callee_argument_dimension(knl, 'callee_fn2')

knl = lp.inline_callable_kernel(knl, 'callee_fn1')

knl = lp.set_options(knl, "write_cl")
knl = lp.set_options(knl, "return_dict")
evt, out_dict = knl(queue, x1=x1, x2=x2)

y1 = out_dict['y1'].get()
y2 = out_dict['y2'].get()

print(la.norm(y1-2*x1.get()))
print(la.norm(y2-3*x2.get()))