Skip to content

Port to loopy-with-namedisl #670

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
inducer merged 2 commits into from
Aug 12, 2026
Merged

Port to loopy-with-namedisl #670

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
214a2d6
Port to loopy-with-namedisl
inducer Jul 26, 2026
602b23f
Update baseline
inducer Aug 12, 2026
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
120 changes: 0 additions & 120 deletions .basedpyright/baseline.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4859,14 +4859,6 @@
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 11,
"endColumn": 19,
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
Expand Down Expand Up @@ -6705,14 +6697,6 @@
}
],
"./pytato/target/python/numpy_like.py": [
{
"code": "reportUnusedImport",
"range": {
"startColumn": 15,
"endColumn": 23,
"lineCount": 1
}
},
{
"code": "reportUnannotatedClassAttribute",
"range": {
Expand Down Expand Up @@ -8909,14 +8893,6 @@
}
],
"./test/test_apps.py": [
{
"code": "reportUnusedImport",
"range": {
"startColumn": 42,
"endColumn": 63,
"lineCount": 1
}
},
{
"code": "reportMissingTypeArgument",
"range": {
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -9615,14 +9559,6 @@
"lineCount": 1
}
},
{
"code": "reportUnknownMemberType",
"range": {
"startColumn": 17,
"endColumn": 30,
"lineCount": 1
}
},
{
"code": "reportUnknownMemberType",
"range": {
Expand All @@ -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": {
Expand All @@ -9655,14 +9575,6 @@
"lineCount": 3
}
},
{
"code": "reportUnknownMemberType",
"range": {
"startColumn": 18,
"endColumn": 31,
"lineCount": 4
}
},
{
"code": "reportUnusedParameter",
"range": {
Expand Down Expand Up @@ -10217,14 +10129,6 @@
}
],
"./test/test_distributed.py": [
{
"code": "reportUnusedImport",
"range": {
"startColumn": 42,
"endColumn": 63,
"lineCount": 1
}
},
{
"code": "reportUnknownParameterType",
"range": {
Expand Down Expand Up @@ -10501,14 +10405,6 @@
}
],
"./test/test_linalg.py": [
{
"code": "reportUnusedImport",
"range": {
"startColumn": 42,
"endColumn": 63,
"lineCount": 1
}
},
{
"code": "reportUnknownMemberType",
"range": {
Expand All @@ -10527,14 +10423,6 @@
}
],
"./test/test_pytato.py": [
{
"code": "reportUnusedImport",
"range": {
"startColumn": 42,
"endColumn": 63,
"lineCount": 1
}
},
{
"code": "reportPrivateUsage",
"range": {
Expand Down Expand Up @@ -10647,14 +10535,6 @@
"lineCount": 1
}
},
{
"code": "reportUnknownMemberType",
"range": {
"startColumn": 56,
"endColumn": 69,
"lineCount": 1
}
},
{
"code": "reportUnusedExpression",
"range": {
Expand Down
3 changes: 3 additions & 0 deletions doc/conf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -73,6 +74,8 @@
# pytools
"lp.TemporaryVariable": "class:loopy.TemporaryVariable",
"lp.AddressSpace": "class:loopy.AddressSpace",
# namedisl
"nisl.Set": "class:namedisl.Set",
}


Expand Down
66 changes: 18 additions & 48 deletions pytato/loopy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,21 +25,21 @@
THE SOFTWARE.
"""


import dataclasses
from numbers import Number
from typing import (
TYPE_CHECKING,
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

Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
10 changes: 5 additions & 5 deletions pytato/scalar_expr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 10 additions & 20 deletions pytato/target/loopy/codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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*.
Expand Down Expand Up @@ -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

Expand Down
Loading
Loading