fix-group-size.py

import loopy as lp
import numpy as np
import pyopencl as cl
from loopy.kernel.data import temp_var_scope as scopes

lp.set_caching_enabled(False)

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

def __get_code(k):
    return lp.generate_code(knl)[0]

knl = lp.make_kernel(['{[i]: 0 <= i < 2}'],
        """
                a[i] = i
        """,
        [lp.GlobalArg('a', dtype=np.float64, shape=(2,))],
        target=lp.OpenCLTarget(),
        )

print __get_code(knl)

knl = lp.split_iname(knl, 'i', 4, inner_tag='l.0')

def ggs(insn_ids, ignore_auto):
    grid_size, local_size = orig_knl.get_grid_sizes_for_insn_ids(insn_ids)
    return grid_size, (4,)

orig_knl = knl
knl = knl.copy(overridden_get_grid_sizes_for_insn_ids=ggs)

print __get_code(knl)