diff --git a/doc/reference.rst b/doc/reference.rst index 05ebf4946c8b11a4def69aa7b30e14edda8bd25b..95c783f600366403d7abb128e77910c4fb5f7be7 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -278,6 +278,8 @@ Piecewise Quasi-Affine Expression .. autoclass:: PwAff :members: +.. autofunction:: make_zero_and_vars + Union of Piecewise Quasi-Affine Expressions ------------------------------------------- diff --git a/islpy/__init__.py b/islpy/__init__.py index f9f339e4e89f6af7266a12cbed86eaff5dffab10..1fe2726fa6428e2c0c496a94a0f529d3a6f1668a 100644 --- a/islpy/__init__.py +++ b/islpy/__init__.py @@ -1186,4 +1186,54 @@ def align_two(obj1, obj2, across_dim_types=False): return (obj1, obj2) +def make_zero_and_vars(set_vars, params=[], ctx=None): + """ + :arg set_vars: an iterable of variable names, or a comma-separated string + :arg params: an iterable of variable names, or a comma-separated string + + :return: a dictionary from variable (in *set_vars* and *params*) + to :class:`PwAff` instances that represent each of the + variables. They key '0' is also include and represents + a : + + .. versionadded:: 2016.1.1 + + This function is intended to make it relatively easy to construct sets + programmatically without resorting to string manipulation. + + Usage example:: + + v = isl.make_zero_and_vars("i,j,k", "n") + + myset = ( + v[0].le_set(v["i"] + v["j"]) + & + (v["i"] + v["j"]).lt_set(v["n"]) + & + (v[0].le_set(v["i"])) + & + (v["i"].le_set(13 + v["n"])) + ) + """ + if ctx is None: + ctx = DEFAULT_CONTEXT + + if isinstance(set_vars, str): + set_vars = [s.strip() for s in set_vars.split(",")] + if isinstance(params, str): + params = [s.strip() for s in params.split(",")] + + space = Space.create_from_names(ctx, set=set_vars, params=params) + + result = {} + + zero = Aff.zero_on_domain(LocalSpace.from_space(space)) + result[0] = PwAff.from_aff(zero) + + var_dict = zero.get_var_dict() + for name, (dt, idx) in var_dict.items(): + result[name] = PwAff.from_aff(zero.set_coefficient_val(dt, idx, 1)) + + return result + # vim: foldmethod=marker diff --git a/test/test_isl.py b/test/test_isl.py index 18fa610ba7b96f7ac08291a33581dd6270098510..a414b994ec6be43fa129da3ce58f1c86fafdf6bc 100644 --- a/test/test_isl.py +++ b/test/test_isl.py @@ -254,6 +254,22 @@ def test_codegen(): print(isl_ast_codegen(s)) +def test_set_building(): + v = isl.make_zero_and_vars("i,j,k", "n") + + myset = ( + v[0].le_set(v["i"] + v["j"]) + & + (v["i"] + v["j"]).lt_set(v["n"]) + & + (v[0].le_set(v["i"])) + & + (v["i"].le_set(13 + v["n"])) + ) + + print(myset) + + if __name__ == "__main__": import sys if len(sys.argv) > 1: