diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 33b30b4a9..bf1ad30fd 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -4859,14 +4859,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownVariableType", - "range": { - "startColumn": 11, - "endColumn": 19, - "lineCount": 1 - } - }, { "code": "reportUnknownVariableType", "range": { @@ -6705,14 +6697,6 @@ } ], "./pytato/target/python/numpy_like.py": [ - { - "code": "reportUnusedImport", - "range": { - "startColumn": 15, - "endColumn": 23, - "lineCount": 1 - } - }, { "code": "reportUnannotatedClassAttribute", "range": { @@ -8909,14 +8893,6 @@ } ], "./test/test_apps.py": [ - { - "code": "reportUnusedImport", - "range": { - "startColumn": 42, - "endColumn": 63, - "lineCount": 1 - } - }, { "code": "reportMissingTypeArgument", "range": { @@ -9159,38 +9135,6 @@ } ], "./test/test_codegen.py": [ - { - "code": "reportUnusedImport", - "range": { - "startColumn": 27, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportUnusedImport", - "range": { - "startColumn": 25, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportUnusedImport", - "range": { - "startColumn": 26, - "endColumn": 59, - "lineCount": 1 - } - }, - { - "code": "reportUnusedImport", - "range": { - "startColumn": 42, - "endColumn": 63, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -9615,14 +9559,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 17, - "endColumn": 30, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -9631,22 +9567,6 @@ "lineCount": 3 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 17, - "endColumn": 30, - "lineCount": 4 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 18, - "endColumn": 31, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -9655,14 +9575,6 @@ "lineCount": 3 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 18, - "endColumn": 31, - "lineCount": 4 - } - }, { "code": "reportUnusedParameter", "range": { @@ -10217,14 +10129,6 @@ } ], "./test/test_distributed.py": [ - { - "code": "reportUnusedImport", - "range": { - "startColumn": 42, - "endColumn": 63, - "lineCount": 1 - } - }, { "code": "reportUnknownParameterType", "range": { @@ -10501,14 +10405,6 @@ } ], "./test/test_linalg.py": [ - { - "code": "reportUnusedImport", - "range": { - "startColumn": 42, - "endColumn": 63, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -10527,14 +10423,6 @@ } ], "./test/test_pytato.py": [ - { - "code": "reportUnusedImport", - "range": { - "startColumn": 42, - "endColumn": 63, - "lineCount": 1 - } - }, { "code": "reportPrivateUsage", "range": { @@ -10647,14 +10535,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 56, - "endColumn": 69, - "lineCount": 1 - } - }, { "code": "reportUnusedExpression", "range": { diff --git a/doc/conf.py b/doc/conf.py index ff5ff8b1b..dabf8a38d 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -28,6 +28,7 @@ "loopy": ("https://documen.tician.de/loopy/", None), "sumpy": ("https://documen.tician.de/sumpy/", None), "islpy": ("https://documen.tician.de/islpy/", None), + "namedisl": ("https://documen.tician.de/namedisl/", None), "jax": ("https://docs.jax.dev/en/latest/", None), "mpi4py": ("https://mpi4py.readthedocs.io/en/latest", None), "constantdict": ("https://matthiasdiener.github.io/constantdict/", None), @@ -73,6 +74,8 @@ # pytools "lp.TemporaryVariable": "class:loopy.TemporaryVariable", "lp.AddressSpace": "class:loopy.AddressSpace", + # namedisl + "nisl.Set": "class:namedisl.Set", } diff --git a/pytato/loopy.py b/pytato/loopy.py index 23eeb1eba..1a486561a 100644 --- a/pytato/loopy.py +++ b/pytato/loopy.py @@ -25,7 +25,6 @@ THE SOFTWARE. """ - import dataclasses from numbers import Number from typing import ( @@ -33,13 +32,14 @@ Any, ) -import islpy as isl +import namedisl as nisl import numpy as np from constantdict import constantdict from typing_extensions import override import loopy as lp import pymbolic.primitives as prim +from loopy.symbolic import pw_aff_to_expr from loopy.typing import assert_tuple from pytools import memoize_method @@ -315,29 +315,6 @@ class ShapeInferenceFailure(RuntimeError): # ruff:ignore[error-suffix-on-except pass -def _get_val_in_bset(bset: isl.BasicSet, idim: int) -> ScalarExpression: - """ - Gets the value of *bset*'s *idim*-th set-dim in terms of it's param-dims. - - .. note:: - - Assumes all constraints in *bset* are equality constraints. - """ - from loopy.symbolic import aff_to_expr - - max_val = bset.dim_max(idim) - - assert max_val.is_equal(bset.dim_min(idim)) - - if max_val.n_piece() != 1: - raise NotImplementedError("Shape inference resulted in a piecewise" - " result.") - - (_, aff), = max_val.get_pieces() - - return aff_to_expr(aff) - - def solve_constraints(variables: Sequence[str], parameters: Sequence[str], constraints: Sequence[tuple[ArithmeticExpression, @@ -355,36 +332,29 @@ def solve_constraints(variables: Sequence[str], :returns: A mapping from variable name in *variables* to :class:`ScalarExpression` obtained after solving for them. """ - from loopy.symbolic import aff_from_expr + if not variables: + return {} - space = isl.Space.create_from_names(isl.DEFAULT_CONTEXT, - set=variables, - params=parameters) + from loopy.symbolic import pwaff_from_expr - shape_inference_bset = isl.BasicSet.universe(space) + space = nisl.Space.from_names(out=variables, param=parameters) - for lhs, rhs in constraints: - aff = aff_from_expr(space, lhs-rhs) + shape_inference_set = nisl.Set.universe(space) - shape_inference_bset = (shape_inference_bset - .add_constraint(isl.Constraint - .equality_from_aff(aff))) + v = nisl.pw_affs_from_domain_space(space) + for lhs, rhs in constraints: + shape_inference_set = shape_inference_set & ( + pwaff_from_expr(v, lhs) + .where("==", pwaff_from_expr(v, rhs))) - if shape_inference_bset.is_empty(): + if shape_inference_set.is_empty(): raise ShapeInferenceFailure - solution = {} - - # {{{ get the value of each unknown variable - - for idim in range(shape_inference_bset.dim(isl.dim_type.set)): - arg_name = shape_inference_bset.get_dim_name(isl.dim_type.set, idim) - solved_val = _get_val_in_bset(shape_inference_bset, idim) - solution[arg_name] = solved_val - - # }}} - - return solution + solution = shape_inference_set.as_map([]).as_pw_multi_aff() + return { + name: pw_aff_to_expr(solution[name]) + for name in variables + } # {{{ shape inference helpers diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index 06219e540..1299f1671 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -376,13 +376,13 @@ def get_reduction_induction_variables(expr: Expression) -> AbstractSet[str]: def is_quasi_affine(expr: Expression) -> bool: - import islpy as isl - space = isl.Space.create_from_names( - isl.DEFAULT_CONTEXT, - set=list(get_dependencies(expr)), + import namedisl as nisl + space = nisl.Space.from_names( + param=(), + out=list(get_dependencies(expr)), ) try: - guarded_pwaff_from_expr(space, expr) + guarded_pwaff_from_expr(nisl.pw_affs_from_domain_space(space), expr) except ExpressionToAffineConversionError: return False return True diff --git a/pytato/target/loopy/codegen.py b/pytato/target/loopy/codegen.py index 87f6a1cea..5d85ca62d 100644 --- a/pytato/target/loopy/codegen.py +++ b/pytato/target/loopy/codegen.py @@ -31,7 +31,7 @@ from functools import reduce from typing import TYPE_CHECKING, Any, cast -import islpy as isl +import namedisl as nisl import numpy as np from constantdict import constantdict from typing_extensions import Never, override @@ -1026,7 +1026,7 @@ def shape_to_scalar_expression(shape: ShapeType, def domain_for_shape(dim_names: tuple[str, ...], shape: tuple[ScalarExpression, ...], reductions: dict[str, tuple[ScalarExpression, ScalarExpression]], - ) -> isl.BasicSet: + ) -> nisl.Set: """Create an :class:`islpy.BasicSet` that expresses an appropriate index domain for an array of (potentially symbolic) shape *shape* having reduction dimensions *reductions*. @@ -1059,30 +1059,20 @@ def domain_for_shape(dim_names: tuple[str, ...], param_names = sorted(param_names_set) # Build domain. - dom = isl.BasicSet.universe( - isl.Space.create_from_names(isl.DEFAULT_CONTEXT, - set=set_names, - params=param_names)) + dom = nisl.Set.universe( + nisl.Space.from_names(out=set_names, param=param_names)) # Add constraints. - from loopy.symbolic import aff_from_expr - affs = isl.affs_from_space(dom.space) + from loopy.symbolic import pwaff_from_expr + v = dom.var_pw_affs for iname, dim in zip(dim_names, shape, strict=True): - dom &= affs[0].le_set(affs[iname]) - dom &= affs[iname].lt_set(aff_from_expr(dom.space, dim)) + dom &= v[0].le_set(v[iname]) + dom &= v[iname].where("<", pwaff_from_expr(v, dim)) for iname, (left, right) in reductions.items(): - dom &= aff_from_expr(dom.space, left).to_pw_aff().le_set(affs[iname]) - dom &= affs[iname].lt_set(aff_from_expr(dom.space, right)) - - doms = dom.get_basic_sets() - - if len(doms) == 0: - # empty set - dom = isl.BasicSet.empty(dom.get_space()) - else: - dom, = doms + dom &= pwaff_from_expr(v, left).where("<=", v[iname]) + dom &= v[iname].where("<", pwaff_from_expr(v, right)) return dom diff --git a/requirements.txt b/requirements.txt index b1f1df20a..909770a28 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ git+https://github.com/inducer/pytools.git#egg=pytools git+https://github.com/inducer/pymbolic.git#egg=pymbolic git+https://github.com/inducer/genpy.git#egg=genpy +git+https://github.com/inducer/namedisl.git#egg=namedisl git+https://github.com/inducer/loopy.git#egg=loopy mako diff --git a/test/test_pytato.py b/test/test_pytato.py index c68ab6598..b7a4f8e73 100644 --- a/test/test_pytato.py +++ b/test/test_pytato.py @@ -180,7 +180,7 @@ def test_zero_length_arrays(): assert y.shape == (0, 4) knl = pt.generate_loopy(y).kernel - assert all(dom.is_empty() for dom in knl.domains if dom.total_dim() != 0) + assert all(dom.is_empty() for dom in knl.domains if dom.space.names) def test_concatenate_input_validation():