add tiled matmul example as test; islpy -> namedisl

a-alveyblanc · a-alveyblanc · commit 907873c6d8d8 · 2026-04-06T09:18:05.000-05:00
diff --git a/examples/python/compute-examples/compute-tiled-matmul.py b/examples/python/compute-examples/compute-tiled-matmul.py
@@ -43,11 +43,11 @@ def main(
     # guarantees for the instruction that stores into c
     knl = lp.fix_parameters(knl, M=M, N=N, K=K)
 
+    # shared memory tile-level splitting
     knl = lp.split_iname(knl, "i", bm, inner_iname="ii", outer_iname="io")
     knl = lp.split_iname(knl, "j", bn, inner_iname="ji", outer_iname="jo")
     knl = lp.split_iname(knl, "k", bk, inner_iname="ki", outer_iname="ko")
 
-    # FIXME: Given the input is already tiled, we shouldn't have to supply compute bounds here.
     compute_map_a = nisl.make_map(f"""{{
         [is, ks] -> [ii_s, io, ki_s, ko] :
             is = io * {bm} + ii_s and
@@ -67,7 +67,8 @@ def main(
             compute_map=compute_map_a,
             storage_indices=["ii_s", "ki_s"],
             temporal_inames=["io", "ko", "jo"],
-            temporary_address_space=lp.AddressSpace.LOCAL
+            temporary_address_space=lp.AddressSpace.LOCAL,
+            temporary_dtype=np.float64
         )
 
         knl = compute(
@@ -76,7 +77,8 @@ def main(
             compute_map=compute_map_b,
             storage_indices=["ki_s", "ji_s"],
             temporal_inames=["io", "ko", "jo"],
-            temporary_address_space=lp.AddressSpace.LOCAL
+            temporary_address_space=lp.AddressSpace.LOCAL,
+            temporary_dtype=np.float64
         )
 
     if use_precompute:
@@ -116,7 +118,7 @@ def main(
         print(20*"=", "Tiled matmul report", 20*"=")
         print(f"Problem size:  M = {M:-4},  N = {N:-4},  K = {K:-4}")
         print(f"Tile size   : BM = {bm:-4}, BN = {bn:-4}, BK = {bk:-4}")
-        print(f"Relative error = {la.norm((a @ b) - out) / la.norm(out)}")
+        print(f"Relative error = {la.norm((a @ b) - out) / la.norm(a @ b)}")
         print((40 + len(" Tiled matmul report "))*"=")
 
     if print_device_code:
diff --git a/examples/python/compute-examples/finite-difference-2-5D.py b/examples/python/compute-examples/finite-difference-2-5D.py
@@ -76,11 +76,10 @@ def main(
         compute_map = nisl.make_map(
             f"""
             {{
-                [is, js, ks] -> [io, ii_s, jo, ji_s, k_s] :
-                0 <= ii_s < {bm} and 0 <= ji_s < {bn} and 0 <= k_s < {npts} and
-                is = io * {bm} + ii_s and
-                js = jo * {bn} + ji_s and
-                ks = k_s
+                [is, js, ks] -> [io, ii_s, jo, ji_s, k] :
+                is = io * {bm} + ii_s - {r} and
+                js = jo * {bn} + ji_s - {r} and
+                ks = k
             }}
             """
         )
@@ -89,8 +88,8 @@ def main(
             knl,
             "u_",
             compute_map=compute_map,
-            storage_indices=["ii_s", "ji_s", "k_s"],
-            temporal_inames=["io", "jo"],
+            storage_indices=["ii_s", "ji_s"],
+            temporal_inames=["io", "jo", "k"],
             temporary_name="u_compute",
             temporary_address_space=lp.AddressSpace.LOCAL,
             temporary_dtype=np.float32
diff --git a/loopy/transform/compute.py b/loopy/transform/compute.py
@@ -271,7 +271,7 @@ def compute(
     """
     Inserts an instruction to compute an expression given by :arg:`substitution`
     and replaces all invocations of :arg:`substitution` with the result of the
-    compute instruction.
+    inserted compute instruction.
 
     :arg substitution: The substitution rule for which the compute
     transform should be applied.
@@ -283,8 +283,6 @@ def compute(
     :arg storage_indices: An ordered sequence of names of storage indices. Used
     to create inames for the loops that cover the required set of compute points.
     """
-    # FIXME: use namedisl directly
-    compute_map = compute_map._reconstruct_isl_object()
 
     # {{{ construct necessary pieces; footprint, global usage map
 
@@ -295,125 +293,74 @@ def compute(
         ctx, expander, kernel, substitution, None
     )
 
+    # add compute inames to domain / kernel
+    domain_changer = DomainChanger(kernel, kernel.all_inames())
+    named_domain = nisl.make_basic_set(domain_changer.domain)
+
     _ = expr_gatherer.map_kernel(kernel)
     usage_exprs = expr_gatherer.usage_expressions
 
-    all_exprs = [
-        expr
-        for usage in usage_exprs
-        for expr in usage
-    ]
+    all_exprs = [expr for usage in usage_exprs for expr in usage]
+    usage_inames = set.union(*(gather_vars(expr) for expr in all_exprs))
 
-    space = space_from_exprs(all_exprs)
+    usage_domain = nisl.make_set(f"{{ [{",".join(iname for iname in usage_inames)}] }}")
+    footprint = nisl.make_set(f"{{ [{",".join(idx for idx in storage_indices)}] }}")
 
-    footprint = isl.Set.empty(
-        isl.Space.create_from_names(
-            ctx=space.get_ctx(),
-            set=list(storage_indices)
-        )
+    global_usage_map = nisl.make_map_from_domain_and_range(
+        usage_domain,
+        compute_map.domain()
     )
 
-    # add compute inames to domain / kernel
-    domain_changer = DomainChanger(kernel, kernel.all_inames())
-    domain = domain_changer.domain
-
-    range_space = isl.Space.create_from_names(
-        ctx=space.get_ctx(),
-        set=list(storage_indices)
-    )
-    map_space = space.map_from_domain_and_range(range_space)
-    global_usage_map = isl.Map.empty(map_space)
+    global_usage_map = nisl.make_map(isl.Map.empty(global_usage_map.get_space()))
 
+    usage_substs: Mapping[AccessTuple, nisl.Map] = {}
     for usage in usage_exprs:
 
         # FIXME package sequence of pymbolic exprs -> multipwaff up as a function in loopy.symbolic
-        local_usage_mpwaff = isl.MultiPwAff.zero(map_space)
+        local_usage_mpwaff = isl.MultiPwAff.zero(global_usage_map.get_space())
 
         for i in range(len(storage_indices)):
+            local_space = local_usage_mpwaff.get_at(i).get_space().domain()
             local_usage_mpwaff = local_usage_mpwaff.set_pw_aff(
                 i,
-                pwaff_from_expr(space, usage[i])
+                pwaff_from_expr(local_space, usage[i])
             )
 
-        local_usage_map = local_usage_mpwaff.as_map()
+        local_usage_map = nisl.make_map(local_usage_mpwaff.as_map())
 
-        # FIXME: fix with namedisl
-        # remove unnecessary names from domain and intersect with usage map
-        usage_names = frozenset(
-            local_usage_map.get_dim_name(isl.dim_type.in_, dim)
-            for dim in range(local_usage_map.dim(isl.dim_type.in_))
-        )
-
-        domain_names = frozenset(
-            domain.get_dim_name(isl.dim_type.set, dim)
-            for dim in range(domain.dim(isl.dim_type.set))
-        )
-
-        domain_tmp = domain.project_out_except(
-            usage_names & domain_names, [isl.dim_type.set]
-        )
-
-        local_usage_map = align_map_domain_to_set(local_usage_map, domain_tmp)
-        local_usage_map = local_usage_map.intersect_domain(domain_tmp)
+        local_usage_map = local_usage_map.intersect_domain(named_domain)
         global_usage_map = global_usage_map | local_usage_map
 
-    # {{{ FIXME: this shouldn't need to be done here; will be handled by namedisl
+        local_storage_map = local_usage_map.apply_range(compute_map)
+        relevant_names = gather_vars(usage)
 
-    global_usage_map = global_usage_map.apply_range(compute_map)
-    common_dims = {
-        dim1 : dim2
-        for dim1 in range(global_usage_map.dim(isl.dim_type.in_))
-        for dim2 in range(global_usage_map.dim(isl.dim_type.out))
-        if (
-            global_usage_map.get_dim_name(isl.dim_type.in_, dim1)
-            ==
-            global_usage_map.get_dim_name(isl.dim_type.out, dim2)
+        local_storage_map = local_storage_map.project_out_except(
+            (relevant_names - frozenset(temporal_inames)) | frozenset(storage_indices)
         )
-    }
 
-    for pos1, pos2 in common_dims.items():
-        global_usage_map = global_usage_map.equate(
-            isl.dim_type.in_, pos1,
-            isl.dim_type.out, pos2
-        )
+        usage_substs[tuple(usage)] = local_storage_map
 
-    # }}}
+    global_usage_map = global_usage_map.apply_range(compute_map)
 
     # }}}
 
     # {{{ compute bounds and update kernel domain
 
     footprint = global_usage_map.range()
-    footprint_tmp, domain = isl.align_two(footprint, domain)
-    domain = (domain & footprint_tmp).get_basic_sets()[0]
-
-    new_domains = domain_changer.get_domains_with(domain)
-    kernel = kernel.copy(domains=new_domains)
-
-    # }}}
-
-    # {{{ compute index expressions
-
-    usage_substs: Mapping[AccessTuple, isl.Map] = {}
-    for usage in usage_exprs:
-        # find the relevant names
-        relevant_names = gather_vars(usage)
-
-        # project out irrelevant names
-        relevant_names = frozenset(relevant_names) - frozenset(temporal_inames)
+    footprint = footprint.project_out_except(
+        frozenset(temporal_inames) | frozenset(storage_indices)
+    )
 
-        local_iname_to_storage = global_usage_map.project_out_except(
-            relevant_names,
-            [isl.dim_type.in_]
-        )
+    # FIXME: probably do not want this permanently here
+    footprint = nisl.make_set(footprint._reconstruct_isl_object().convex_hull())
+    named_domain = named_domain & footprint
 
-        local_iname_to_storage = local_iname_to_storage.project_out_except(
-            storage_indices,
-            [isl.dim_type.out]
-        )
+    # FIXME:
+    if len(named_domain.get_basic_sets()) != 1:
+        raise ValueError("New domain should be composed of a single basic set")
 
-        # map usage -> resulting map
-        usage_substs[tuple(usage)] = local_iname_to_storage
+    new_domains = domain_changer.get_domains_with(named_domain.get_basic_sets()[0])
+    kernel = kernel.copy(domains=new_domains)
 
     # }}}
 
@@ -487,9 +434,7 @@ def compute(
     loopy_type = to_loopy_type(temporary_dtype, allow_none=True)
 
     # FIXME: fix with namedisl?
-    shape_domain = footprint.project_out_except(storage_indices,
-                                                [isl.dim_type.set])
-    shape_domain = shape_domain.project_out_except("", [isl.dim_type.param])
+    shape_domain = footprint.project_out_except(storage_indices)
 
     temp_shape = tuple(
         pw_aff_to_expr(shape_domain.dim_max(dim)) + 1
diff --git a/test/test_transform.py b/test/test_transform.py
@@ -20,6 +20,7 @@
 THE SOFTWARE.
 """
 
+from collections.abc import Mapping
 import logging
 
 import numpy as np
@@ -1745,6 +1746,108 @@ def test_duplicate_iname_not_read_only_nested(ctx_factory: cl.CtxFactory):
     lp.auto_test_vs_ref(ref_t_unit, ctx, t_unit)
 
 
+@pytest.mark.parametrize("case", (
+    {"M": 128, "N": 128, "K": 128, "BM": 32, "BN": 32, "BK": 16},
+    {"M": 200, "N": 200, "K": 200, "BM": 32, "BN": 32, "BK": 16},
+))
+def test_compute_simple_tiled_matmul(
+        ctx_factory: cl.CtxFactory,
+        case: Mapping[str, int]
+    ):
+
+    import namedisl as nisl
+
+    M = case["M"]
+    N = case["N"]
+    K = case["K"]
+    bm = case["BM"]
+    bn = case["BN"]
+    bk = case["BK"]
+
+    knl = lp.make_kernel(
+        "{ [i, j, k] : 0 <= i < M and 0 <= j < N and 0 <= k < K }",
+        """
+        a_(is, ks) := a[is, ks]
+        b_(ks, js) := b[ks, js]
+        c[i, j] = sum([k], a_(i, k) * b_(k, j))
+        """,
+        [
+            lp.GlobalArg("a", shape=(M, K), dtype=np.float64),
+            lp.GlobalArg("b", shape=(K, N), dtype=np.float64),
+            lp.GlobalArg("c", shape=(M, N), dtype=np.float64,
+                         is_output=True)
+        ]
+    )
+
+    knl = lp.fix_parameters(knl, M=M, N=N, K=K)
+
+    # shared memory tile-level splitting
+    knl = lp.split_iname(knl, "i", bm, inner_iname="ii", outer_iname="io")
+    knl = lp.split_iname(knl, "j", bn, inner_iname="ji", outer_iname="jo")
+    knl = lp.split_iname(knl, "k", bk, inner_iname="ki", outer_iname="ko")
+
+    compute_map_a = nisl.make_map(f"""{{
+        [is, ks] -> [ii_s, io, ki_s, ko] :
+            is = io * {bm} + ii_s and
+            ks = ko * {bk} + ki_s
+    }}""")
+
+    compute_map_b = nisl.make_map(f"""{{
+        [ks, js] -> [ki_s, ko, ji_s, jo] :
+            js = jo * {bn} + ji_s and
+            ks = ko * {bk} + ki_s
+    }}""")
+
+    from loopy.transform.compute import compute
+    knl = compute(
+        knl,
+        "a_",
+        compute_map=compute_map_a,
+        storage_indices=["ii_s", "ki_s"],
+        temporal_inames=["io", "ko", "jo"],
+        temporary_address_space=lp.AddressSpace.LOCAL,
+        temporary_dtype=np.float64
+    )
+
+    knl = compute(
+        knl,
+        "b_",
+        compute_map=compute_map_b,
+        storage_indices=["ki_s", "ji_s"],
+        temporal_inames=["io", "ko", "jo"],
+        temporary_address_space=lp.AddressSpace.LOCAL,
+        temporary_dtype=np.float64
+    )
+
+    knl = lp.tag_inames(
+        knl, {
+            "io"   : "g.0", # outer block loop over block rows
+            "jo"   : "g.1", # outer block loop over block cols
+
+            "ii"   : "l.0", # inner block loop over rows
+            "ji"   : "l.1", # inner block loop over cols
+
+            "ii_s" : "l.0", # inner storage loop over a rows
+            "ji_s" : "l.0", # inner storage loop over b cols
+            "ki_s" : "l.1"  # inner storage loop over a cols / b rows
+        }
+    )
+
+    knl = lp.add_inames_for_unused_hw_axes(knl)
+
+    ctx = ctx_factory()
+    queue = cl.CommandQueue(ctx)
+
+    a = np.random.randn(M, K)
+    b = np.random.randn(K, N)
+
+    ex = knl.executor(ctx)
+    _, out = ex(queue, a=a, b=b)
+
+    import numpy.linalg as la
+    assert (la.norm((a @ b) - out) / la.norm(a @ b)) < 1e-15
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1:
