quack-kernels 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

quack/sort/generate_sorting_networks.py

@@ -0,0 +1,326 @@
+#!/usr/bin/env python3
+"""
+Generate optimized sorting network code from the optimal sorting network data.
+Based on data from: https://bertdobbelaere.github.io/sorting_networks.html
+
+This script generates CUTE DSL functions for optimal sorting networks of various sizes.
+"""
+
+import argparse
+import os
+import re
+from typing import List, Tuple, Dict
+
+# Network strings from bertdobbelaere.github.io/sorting_networks.html
+# Copy-paste network strings here, then run initialize_networks() to parse them
+NETWORK_STRINGS = {
+    # Size 2: 1 CE, depth 1
+    2: """
+[(0,1)]
+    """,
+    # Size 4: 5 CEs, depth 3
+    4: """
+[(0,2),(1,3)]
+[(0,1),(2,3)]
+[(1,2)]
+    """,
+    # Size 8: 19 CEs, depth 6
+    8: """
+[(0,2),(1,3),(4,6),(5,7)]
+[(0,4),(1,5),(2,6),(3,7)]
+[(0,1),(2,3),(4,5),(6,7)]
+[(2,4),(3,5)]
+[(1,4),(3,6)]
+[(1,2),(3,4),(5,6)]
+    """,
+    # Size 16: 60 CEs, depth 10
+    16: """
+[(0,13),(1,12),(2,15),(3,14),(4,8),(5,6),(7,11),(9,10)]
+[(0,5),(1,7),(2,9),(3,4),(6,13),(8,14),(10,15),(11,12)]
+[(0,1),(2,3),(4,5),(6,8),(7,9),(10,11),(12,13),(14,15)]
+[(0,2),(1,3),(4,10),(5,11),(6,7),(8,9),(12,14),(13,15)]
+[(1,2),(3,12),(4,6),(5,7),(8,10),(9,11),(13,14)]
+[(1,4),(2,6),(5,8),(7,10),(9,13),(11,14)]
+[(2,4),(3,6),(9,12),(11,13)]
+[(3,5),(6,8),(7,9),(10,12)]
+[(3,4),(5,6),(7,8),(9,10),(11,12)]
+[(6,7),(8,9)]
+    """,
+    # Size 32: 185 CEs, depth 14
+    32: """
+[(0,1),(2,3),(4,5),(6,7),(8,9),(10,11),(12,13),(14,15),(16,17),(18,19),(20,21),(22,23),(24,25),(26,27),(28,29),(30,31)]
+[(0,2),(1,3),(4,6),(5,7),(8,10),(9,11),(12,14),(13,15),(16,18),(17,19),(20,22),(21,23),(24,26),(25,27),(28,30),(29,31)]
+[(0,4),(1,5),(2,6),(3,7),(8,12),(9,13),(10,14),(11,15),(16,20),(17,21),(18,22),(19,23),(24,28),(25,29),(26,30),(27,31)]
+[(0,8),(1,9),(2,10),(3,11),(4,12),(5,13),(6,14),(7,15),(16,24),(17,25),(18,26),(19,27),(20,28),(21,29),(22,30),(23,31)]
+[(0,16),(1,8),(2,4),(3,12),(5,10),(6,9),(7,14),(11,13),(15,31),(17,24),(18,20),(19,28),(21,26),(22,25),(23,30),(27,29)]
+[(1,2),(3,5),(4,8),(6,22),(7,11),(9,25),(10,12),(13,14),(17,18),(19,21),(20,24),(23,27),(26,28),(29,30)]
+[(1,17),(2,18),(3,19),(4,20),(5,10),(7,23),(8,24),(11,27),(12,28),(13,29),(14,30),(21,26)]
+[(3,17),(4,16),(5,21),(6,18),(7,9),(8,20),(10,26),(11,23),(13,25),(14,28),(15,27),(22,24)]
+[(1,4),(3,8),(5,16),(7,17),(9,21),(10,22),(11,19),(12,20),(14,24),(15,26),(23,28),(27,30)]
+[(2,5),(7,8),(9,18),(11,17),(12,16),(13,22),(14,20),(15,19),(23,24),(26,29)]
+[(2,4),(6,12),(9,16),(10,11),(13,17),(14,18),(15,22),(19,25),(20,21),(27,29)]
+[(5,6),(8,12),(9,10),(11,13),(14,16),(15,17),(18,20),(19,23),(21,22),(25,26)]
+[(3,5),(6,7),(8,9),(10,12),(11,14),(13,16),(15,18),(17,20),(19,21),(22,23),(24,25),(26,28)]
+[(3,4),(5,6),(7,8),(9,10),(11,12),(13,14),(15,16),(17,18),(19,20),(21,22),(23,24),(25,26),(27,28)]
+    """,
+    # Size 64: 512 CEs, depth 21
+    64: """
+[(0,2),(1,3),(4,6),(5,7),(8,10),(9,11),(12,14),(13,15),(16,18),(17,19),(20,22),(21,23),(24,26),(25,27),(28,30),(29,31),(32,34),(33,35),(36,38),(37,39),(40,42),(41,43),(44,46),(45,47),(48,50),(49,51),(52,54),(53,55),(56,58),(57,59),(60,62),(61,63)]
+[(0,1),(2,3),(4,5),(6,7),(8,9),(10,11),(12,13),(14,15),(16,17),(18,19),(20,21),(22,23),(24,25),(26,27),(28,29),(30,31),(32,33),(34,35),(36,37),(38,39),(40,41),(42,43),(44,45),(46,47),(48,49),(50,51),(52,53),(54,55),(56,57),(58,59),(60,61),(62,63)]
+[(0,52),(1,2),(3,55),(4,48),(5,6),(7,51),(8,60),(9,10),(11,63),(12,56),(13,14),(15,59),(16,32),(17,18),(19,35),(20,24),(21,22),(23,27),(25,26),(28,44),(29,30),(31,47),(33,34),(36,40),(37,38),(39,43),(41,42),(45,46),(49,50),(53,54),(57,58),(61,62)]
+[(0,20),(1,53),(2,54),(3,23),(4,28),(5,49),(6,50),(7,31),(8,36),(9,61),(10,62),(11,39),(12,16),(13,57),(14,58),(15,19),(17,33),(18,34),(21,25),(22,26),(24,52),(27,55),(29,45),(30,46),(32,56),(35,59),(37,41),(38,42),(40,60),(43,63),(44,48),(47,51)]
+[(0,4),(1,21),(2,22),(3,7),(5,29),(6,30),(8,12),(9,37),(10,38),(11,15),(13,17),(14,18),(16,20),(19,23),(24,32),(25,53),(26,54),(27,35),(28,36),(31,39),(33,57),(34,58),(40,44),(41,61),(42,62),(43,47),(45,49),(46,50),(48,52),(51,55),(56,60),(59,63)]
+[(0,8),(1,5),(2,6),(3,11),(4,12),(7,15),(9,13),(10,14),(16,40),(17,21),(18,22),(19,43),(20,44),(23,47),(24,28),(25,33),(26,34),(27,31),(29,37),(30,38),(32,36),(35,39),(41,45),(42,46),(48,56),(49,53),(50,54),(51,59),(52,60),(55,63),(57,61),(58,62)]
+[(1,9),(2,10),(4,8),(5,13),(6,14),(7,11),(12,48),(15,51),(16,24),(17,41),(18,42),(19,27),(20,28),(21,45),(22,46),(23,31),(25,29),(26,30),(32,40),(33,37),(34,38),(35,43),(36,44),(39,47),(49,57),(50,58),(52,56),(53,61),(54,62),(55,59)]
+[(4,16),(5,9),(6,10),(7,19),(8,24),(11,27),(13,49),(14,50),(17,25),(18,26),(20,32),(21,29),(22,30),(23,35),(28,40),(31,43),(33,41),(34,42),(36,52),(37,45),(38,46),(39,55),(44,56),(47,59),(53,57),(54,58)]
+[(1,4),(5,17),(6,18),(8,16),(9,25),(10,26),(11,19),(12,24),(15,27),(21,33),(22,34),(29,41),(30,42),(36,48),(37,53),(38,54),(39,51),(44,52),(45,57),(46,58),(47,55),(59,62)]
+[(2,8),(9,17),(10,18),(12,20),(13,25),(14,26),(15,23),(24,32),(27,35),(28,36),(31,39),(37,49),(38,50),(40,48),(43,51),(45,53),(46,54),(55,61)]
+[(2,4),(12,16),(13,21),(14,22),(15,19),(20,24),(23,27),(25,33),(26,34),(28,32),(29,37),(30,38),(31,35),(36,40),(39,43),(41,49),(42,50),(44,48),(47,51),(59,61)]
+[(4,16),(5,20),(10,40),(13,17),(14,18),(21,25),(22,26),(23,53),(24,28),(27,31),(29,33),(30,34),(32,36),(35,39),(37,41),(38,42),(43,58),(45,49),(46,50),(47,59)]
+[(3,17),(6,36),(7,21),(8,32),(9,24),(11,41),(13,28),(14,44),(15,45),(18,48),(19,49),(22,52),(25,29),(26,30),(27,57),(31,55),(33,37),(34,38),(35,50),(39,54),(42,56),(46,60)]
+[(6,20),(8,16),(10,24),(11,25),(14,28),(15,29),(17,33),(18,32),(21,37),(22,36),(26,42),(27,41),(30,46),(31,45),(34,48),(35,49),(38,52),(39,53),(43,57),(47,55)]
+[(3,18),(5,8),(6,12),(7,22),(15,21),(17,32),(19,33),(23,37),(26,40),(30,44),(31,46),(41,56),(42,48),(45,60),(51,57),(55,58)]
+[(3,16),(7,20),(11,26),(18,24),(19,25),(22,28),(23,29),(27,33),(30,36),(34,40),(35,41),(37,52),(38,44),(39,45),(43,56),(47,60)]
+[(3,9),(7,13),(10,16),(11,17),(14,20),(15,30),(19,34),(21,36),(23,38),(25,40),(26,32),(27,42),(29,44),(31,37),(33,48),(43,49),(46,52),(47,53),(50,56),(54,60)]
+[(3,8),(7,10),(9,12),(11,18),(13,14),(15,24),(17,22),(19,28),(21,26),(23,25),(27,34),(29,36),(30,32),(31,33),(35,44),(37,42),(38,40),(39,48),(41,46),(45,52),(49,50),(51,54),(53,56),(55,60)]
+[(3,6),(7,12),(11,16),(15,17),(18,20),(19,24),(21,22),(23,30),(25,32),(26,28),(27,29),(31,38),(33,40),(34,36),(35,37),(39,44),(41,42),(43,45),(46,48),(47,52),(51,56),(57,60)]
+[(3,5),(6,8),(7,9),(10,12),(11,13),(14,16),(15,18),(17,20),(19,21),(22,24),(23,26),(25,28),(27,30),(29,32),(31,34),(33,36),(35,38),(37,40),(39,41),(42,44),(43,46),(45,48),(47,49),(50,52),(51,53),(54,56),(55,57),(58,60)]
+[(3,4),(7,8),(11,12),(13,14),(15,16),(17,18),(19,20),(21,22),(23,24),(25,26),(27,28),(29,30),(31,32),(33,34),(35,36),(37,38),(39,40),(41,42),(43,44),(45,46),(47,48),(49,50),(51,52),(55,56),(59,60)]
+    """,
+}
+
+# This will be populated by initialize_networks()
+OPTIMAL_NETWORKS: Dict[int, Tuple[int, int, List[List[Tuple[int, int]]]]] = {}
+
+
+def parse_network_string(network_str: str) -> List[List[Tuple[int, int]]]:
+    """
+    Parse a sorting network string from bertdobbelaere.github.io format.
+
+    Examples:
+    Input: "[(0,2),(1,3)], [(0,1),(2,3)], [(1,2)]"
+    Output: [[(0, 2), (1, 3)], [(0, 1), (2, 3)], [(1, 2)]]
+
+    Input: "[(0,1)], [(1,2)], [(0,1)]"
+    Output: [[(0, 1)], [(1, 2)], [(0, 1)]]
+    """
+    # Remove whitespace and split by '], ['
+    network_str = network_str.strip()
+    if not network_str:
+        return []
+
+    # Split into layer strings
+    layer_pattern = r"\[((?:\(\d+,\d+\)(?:,\(\d+,\d+\))*)?)\]"
+    layers = []
+
+    for match in re.finditer(layer_pattern, network_str):
+        layer_str = match.group(1)
+        if not layer_str.strip():
+            layers.append([])
+            continue
+
+        # Parse comparisons in this layer: (i,j), (k,l), ...
+        comparisons = []
+        comp_pattern = r"\((\d+),(\d+)\)"
+
+        for comp_match in re.finditer(comp_pattern, layer_str):
+            i, j = int(comp_match.group(1)), int(comp_match.group(2))
+            comparisons.append((i, j))
+
+        layers.append(comparisons)
+
+    return layers
+
+
+def calculate_network_stats(layers: List[List[Tuple[int, int]]]) -> Tuple[int, int, int]:
+    """Calculate depth, total comparisons, and max index from network layers."""
+    depth = len(layers)
+    total_comparisons = sum(len(layer) for layer in layers)
+
+    # Find maximum index to determine network size
+    max_index = 0
+    for layer in layers:
+        for i, j in layer:
+            max_index = max(max_index, i, j)
+
+    network_size = max_index + 1  # Since indices are 0-based
+    return depth, total_comparisons, network_size
+
+
+def add_network_from_string(size: int, network_str: str, description: str = ""):
+    """
+    Add a network from a string representation to the OPTIMAL_NETWORKS dictionary.
+
+    Args:
+        size: Size of the network (number of elements)
+        network_str: Network string in bertdobbelaere.github.io format
+        description: Optional description for debugging
+    """
+    try:
+        layers = parse_network_string(network_str)
+        depth, comparisons, detected_size = calculate_network_stats(layers)
+
+        if detected_size != size:
+            print(f"Warning: Network size mismatch! Expected {size}, detected {detected_size}")
+            print(f"Network string: {network_str[:100]}...")
+            return False
+
+        OPTIMAL_NETWORKS[size] = (depth, comparisons, layers)
+
+        if description:
+            print(f"Added network for size {size}: {description}")
+        print(f"  Depth: {depth}, Comparisons: {comparisons}")
+        return True
+
+    except Exception as e:
+        print(f"Error parsing network for size {size}: {e}")
+        print(f"Network string: {network_str[:100]}...")
+        return False
+
+
+def generate_networks_dict(
+    networks_data: Dict[int, Tuple[int, int, List[List[Tuple[int, int]]]]]
+) -> str:
+    """Generate the global networks dictionary."""
+    lines = ["networks = {"]
+
+    for size, (depth, num_comparisons, layers) in sorted(networks_data.items()):
+        # Format the network with proper indentation and newlines
+        network_lines = []
+        for i, layer in enumerate(layers):
+            if i == 0:
+                network_lines.append(f"            {layer}")
+            else:
+                network_lines.append(f",\n            {layer}")
+
+        if len(layers) == 1:
+            network_str = f"[{network_lines[0].strip()}]"
+        else:
+            network_str = "[\n" + "".join(network_lines) + "\n        ]"
+
+        lines.append(f"    # Size {size}: {num_comparisons} CEs, depth {depth}")
+        lines.append(f"    {size}: {network_str},")
+        lines.append("")
+
+    lines.append("}")
+    return "\n".join(lines)
+
+
+def generate_optimal_sort_function() -> str:
+    """Generate the single optimal_sort function that looks up networks by size."""
+    return """@cute.jit
+def optimal_sort(
+    arr: cute.Tensor,
+    n: cutlass.Constexpr[int],
+    start: cutlass.Constexpr[int] = 0,
+    ascending: cutlass.Constexpr[bool] = True
+) -> None:
+    \"\"\"
+    Optimal sorting network dispatcher.
+
+    Args:
+        arr: Array to sort
+        n: Size of array (must be power of 2 and available in networks)
+        start: Starting index (default 0)
+        ascending: Sort in ascending order (default True)
+
+    Source: https://bertdobbelaere.github.io/sorting_networks.html
+    \"\"\"
+    assert n in networks
+    for level in networks[n]:
+        for i, j in level:
+            compare_and_swap(arr, start + i, start + j, ascending)
+"""
+
+
+def generate_sorting_networks_file(max_size: int = 64):
+    """Generate a complete sorting networks file with optimal networks up to max_size."""
+
+    output_file = os.path.join(os.path.dirname(__file__), "sorting_networks.py")
+
+    # Header
+    header = '''# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Tri Dao.
+"""
+Optimal sorting networks generated from: https://bertdobbelaere.github.io/sorting_networks.html
+
+This file was auto-generated by quack/sort/generate_sorting_networks.py. Do not edit it directly.
+"""
+
+# fmt: off
+# ruff: noqa
+# isort: skip_file
+
+import cutlass
+import cutlass.cute as cute
+
+from quack.sort.utils import compare_and_swap
+
+
+'''
+
+    # Generate networks dictionary and optimal_sort function
+    sizes = [n for n in range(2, max_size + 1) if n in OPTIMAL_NETWORKS]
+    networks_dict = generate_networks_dict(OPTIMAL_NETWORKS)
+    optimal_sort_func = generate_optimal_sort_function()
+
+    # Combine everything
+    content = header + networks_dict + "\n\n\n" + optimal_sort_func
+
+    with open(output_file, "w") as f:
+        f.write(content)
+
+    print(f"Generated optimal sorting networks for sizes {sizes}")
+    print(f"Output written to: {output_file}")
+    return sizes
+
+
+def initialize_networks():
+    """Initialize the OPTIMAL_NETWORKS dictionary by parsing NETWORK_STRINGS."""
+    global OPTIMAL_NETWORKS
+    OPTIMAL_NETWORKS.clear()
+
+    for size, network_str in NETWORK_STRINGS.items():
+        success = add_network_from_string(size, network_str, f"Size {size} optimal network")
+        if not success:
+            print(f"Warning: Failed to parse network for size {size}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate optimal sorting network code from bertdobbelaere.github.io data"
+    )
+    parser.add_argument(
+        "--max-size",
+        "-m",
+        type=int,
+        default=64,
+        help="Maximum sorting network size to generate (default: 32)",
+    )
+    parser.add_argument(
+        "--stats", "-s", action="store_true", help="Print statistics about the optimal networks"
+    )
+
+    args = parser.parse_args()
+
+    # Initialize networks from strings
+    initialize_networks()
+
+    if args.stats:
+        print("Optimal Sorting Network Statistics:")
+        print("Size\tDepth\tComparisons\tLayers")
+        print("-" * 35)
+        for n in sorted(OPTIMAL_NETWORKS.keys()):
+            if n <= args.max_size:
+                depth, comparisons, layers = OPTIMAL_NETWORKS[n]
+                print(f"{n}\t{depth}\t{comparisons}\t\t{len(layers)}")
+
+    # Generate the sorting networks file
+    sizes = generate_sorting_networks_file(args.max_size)
+
+    print(f"\nGenerated optimal sorting networks for {len(sizes)} sizes")
+    print(f"Total networks: {len(sizes)}")
+    print(f"Max network size: {max(sizes)}")
+
+
+if __name__ == "__main__":
+    main()

quack/sort/sorting_networks.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Tri Dao.
+"""
+Optimal sorting networks generated from: https://bertdobbelaere.github.io/sorting_networks.html
+
+This file was auto-generated by quack/sort/generate_sorting_networks.py. Do not edit it directly.
+"""
+
+# fmt: off
+# ruff: noqa
+# isort: skip_file
+
+import cutlass
+import cutlass.cute as cute
+
+from quack.sort.utils import compare_and_swap
+
+
+networks = {
+    # Size 2: 1 CEs, depth 1
+    2: [[(0, 1)]],
+
+    # Size 4: 5 CEs, depth 3
+    4: [
+            [(0, 2), (1, 3)],
+            [(0, 1), (2, 3)],
+            [(1, 2)]
+        ],
+
+    # Size 8: 19 CEs, depth 6
+    8: [
+            [(0, 2), (1, 3), (4, 6), (5, 7)],
+            [(0, 4), (1, 5), (2, 6), (3, 7)],
+            [(0, 1), (2, 3), (4, 5), (6, 7)],
+            [(2, 4), (3, 5)],
+            [(1, 4), (3, 6)],
+            [(1, 2), (3, 4), (5, 6)]
+        ],
+
+    # Size 16: 60 CEs, depth 10
+    16: [
+            [(0, 13), (1, 12), (2, 15), (3, 14), (4, 8), (5, 6), (7, 11), (9, 10)],
+            [(0, 5), (1, 7), (2, 9), (3, 4), (6, 13), (8, 14), (10, 15), (11, 12)],
+            [(0, 1), (2, 3), (4, 5), (6, 8), (7, 9), (10, 11), (12, 13), (14, 15)],
+            [(0, 2), (1, 3), (4, 10), (5, 11), (6, 7), (8, 9), (12, 14), (13, 15)],
+            [(1, 2), (3, 12), (4, 6), (5, 7), (8, 10), (9, 11), (13, 14)],
+            [(1, 4), (2, 6), (5, 8), (7, 10), (9, 13), (11, 14)],
+            [(2, 4), (3, 6), (9, 12), (11, 13)],
+            [(3, 5), (6, 8), (7, 9), (10, 12)],
+            [(3, 4), (5, 6), (7, 8), (9, 10), (11, 12)],
+            [(6, 7), (8, 9)]
+        ],
+
+    # Size 32: 185 CEs, depth 14
+    32: [
+            [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9), (10, 11), (12, 13), (14, 15), (16, 17), (18, 19), (20, 21), (22, 23), (24, 25), (26, 27), (28, 29), (30, 31)],
+            [(0, 2), (1, 3), (4, 6), (5, 7), (8, 10), (9, 11), (12, 14), (13, 15), (16, 18), (17, 19), (20, 22), (21, 23), (24, 26), (25, 27), (28, 30), (29, 31)],
+            [(0, 4), (1, 5), (2, 6), (3, 7), (8, 12), (9, 13), (10, 14), (11, 15), (16, 20), (17, 21), (18, 22), (19, 23), (24, 28), (25, 29), (26, 30), (27, 31)],
+            [(0, 8), (1, 9), (2, 10), (3, 11), (4, 12), (5, 13), (6, 14), (7, 15), (16, 24), (17, 25), (18, 26), (19, 27), (20, 28), (21, 29), (22, 30), (23, 31)],
+            [(0, 16), (1, 8), (2, 4), (3, 12), (5, 10), (6, 9), (7, 14), (11, 13), (15, 31), (17, 24), (18, 20), (19, 28), (21, 26), (22, 25), (23, 30), (27, 29)],
+            [(1, 2), (3, 5), (4, 8), (6, 22), (7, 11), (9, 25), (10, 12), (13, 14), (17, 18), (19, 21), (20, 24), (23, 27), (26, 28), (29, 30)],
+            [(1, 17), (2, 18), (3, 19), (4, 20), (5, 10), (7, 23), (8, 24), (11, 27), (12, 28), (13, 29), (14, 30), (21, 26)],
+            [(3, 17), (4, 16), (5, 21), (6, 18), (7, 9), (8, 20), (10, 26), (11, 23), (13, 25), (14, 28), (15, 27), (22, 24)],
+            [(1, 4), (3, 8), (5, 16), (7, 17), (9, 21), (10, 22), (11, 19), (12, 20), (14, 24), (15, 26), (23, 28), (27, 30)],
+            [(2, 5), (7, 8), (9, 18), (11, 17), (12, 16), (13, 22), (14, 20), (15, 19), (23, 24), (26, 29)],
+            [(2, 4), (6, 12), (9, 16), (10, 11), (13, 17), (14, 18), (15, 22), (19, 25), (20, 21), (27, 29)],
+            [(5, 6), (8, 12), (9, 10), (11, 13), (14, 16), (15, 17), (18, 20), (19, 23), (21, 22), (25, 26)],
+            [(3, 5), (6, 7), (8, 9), (10, 12), (11, 14), (13, 16), (15, 18), (17, 20), (19, 21), (22, 23), (24, 25), (26, 28)],
+            [(3, 4), (5, 6), (7, 8), (9, 10), (11, 12), (13, 14), (15, 16), (17, 18), (19, 20), (21, 22), (23, 24), (25, 26), (27, 28)]
+        ],
+
+    # Size 64: 521 CEs, depth 21
+    64: [
+            [(0, 2), (1, 3), (4, 6), (5, 7), (8, 10), (9, 11), (12, 14), (13, 15), (16, 18), (17, 19), (20, 22), (21, 23), (24, 26), (25, 27), (28, 30), (29, 31), (32, 34), (33, 35), (36, 38), (37, 39), (40, 42), (41, 43), (44, 46), (45, 47), (48, 50), (49, 51), (52, 54), (53, 55), (56, 58), (57, 59), (60, 62), (61, 63)],
+            [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9), (10, 11), (12, 13), (14, 15), (16, 17), (18, 19), (20, 21), (22, 23), (24, 25), (26, 27), (28, 29), (30, 31), (32, 33), (34, 35), (36, 37), (38, 39), (40, 41), (42, 43), (44, 45), (46, 47), (48, 49), (50, 51), (52, 53), (54, 55), (56, 57), (58, 59), (60, 61), (62, 63)],
+            [(0, 52), (1, 2), (3, 55), (4, 48), (5, 6), (7, 51), (8, 60), (9, 10), (11, 63), (12, 56), (13, 14), (15, 59), (16, 32), (17, 18), (19, 35), (20, 24), (21, 22), (23, 27), (25, 26), (28, 44), (29, 30), (31, 47), (33, 34), (36, 40), (37, 38), (39, 43), (41, 42), (45, 46), (49, 50), (53, 54), (57, 58), (61, 62)],
+            [(0, 20), (1, 53), (2, 54), (3, 23), (4, 28), (5, 49), (6, 50), (7, 31), (8, 36), (9, 61), (10, 62), (11, 39), (12, 16), (13, 57), (14, 58), (15, 19), (17, 33), (18, 34), (21, 25), (22, 26), (24, 52), (27, 55), (29, 45), (30, 46), (32, 56), (35, 59), (37, 41), (38, 42), (40, 60), (43, 63), (44, 48), (47, 51)],
+            [(0, 4), (1, 21), (2, 22), (3, 7), (5, 29), (6, 30), (8, 12), (9, 37), (10, 38), (11, 15), (13, 17), (14, 18), (16, 20), (19, 23), (24, 32), (25, 53), (26, 54), (27, 35), (28, 36), (31, 39), (33, 57), (34, 58), (40, 44), (41, 61), (42, 62), (43, 47), (45, 49), (46, 50), (48, 52), (51, 55), (56, 60), (59, 63)],
+            [(0, 8), (1, 5), (2, 6), (3, 11), (4, 12), (7, 15), (9, 13), (10, 14), (16, 40), (17, 21), (18, 22), (19, 43), (20, 44), (23, 47), (24, 28), (25, 33), (26, 34), (27, 31), (29, 37), (30, 38), (32, 36), (35, 39), (41, 45), (42, 46), (48, 56), (49, 53), (50, 54), (51, 59), (52, 60), (55, 63), (57, 61), (58, 62)],
+            [(1, 9), (2, 10), (4, 8), (5, 13), (6, 14), (7, 11), (12, 48), (15, 51), (16, 24), (17, 41), (18, 42), (19, 27), (20, 28), (21, 45), (22, 46), (23, 31), (25, 29), (26, 30), (32, 40), (33, 37), (34, 38), (35, 43), (36, 44), (39, 47), (49, 57), (50, 58), (52, 56), (53, 61), (54, 62), (55, 59)],
+            [(4, 16), (5, 9), (6, 10), (7, 19), (8, 24), (11, 27), (13, 49), (14, 50), (17, 25), (18, 26), (20, 32), (21, 29), (22, 30), (23, 35), (28, 40), (31, 43), (33, 41), (34, 42), (36, 52), (37, 45), (38, 46), (39, 55), (44, 56), (47, 59), (53, 57), (54, 58)],
+            [(1, 4), (5, 17), (6, 18), (8, 16), (9, 25), (10, 26), (11, 19), (12, 24), (15, 27), (21, 33), (22, 34), (29, 41), (30, 42), (36, 48), (37, 53), (38, 54), (39, 51), (44, 52), (45, 57), (46, 58), (47, 55), (59, 62)],
+            [(2, 8), (9, 17), (10, 18), (12, 20), (13, 25), (14, 26), (15, 23), (24, 32), (27, 35), (28, 36), (31, 39), (37, 49), (38, 50), (40, 48), (43, 51), (45, 53), (46, 54), (55, 61)],
+            [(2, 4), (12, 16), (13, 21), (14, 22), (15, 19), (20, 24), (23, 27), (25, 33), (26, 34), (28, 32), (29, 37), (30, 38), (31, 35), (36, 40), (39, 43), (41, 49), (42, 50), (44, 48), (47, 51), (59, 61)],
+            [(4, 16), (5, 20), (10, 40), (13, 17), (14, 18), (21, 25), (22, 26), (23, 53), (24, 28), (27, 31), (29, 33), (30, 34), (32, 36), (35, 39), (37, 41), (38, 42), (43, 58), (45, 49), (46, 50), (47, 59)],
+            [(3, 17), (6, 36), (7, 21), (8, 32), (9, 24), (11, 41), (13, 28), (14, 44), (15, 45), (18, 48), (19, 49), (22, 52), (25, 29), (26, 30), (27, 57), (31, 55), (33, 37), (34, 38), (35, 50), (39, 54), (42, 56), (46, 60)],
+            [(6, 20), (8, 16), (10, 24), (11, 25), (14, 28), (15, 29), (17, 33), (18, 32), (21, 37), (22, 36), (26, 42), (27, 41), (30, 46), (31, 45), (34, 48), (35, 49), (38, 52), (39, 53), (43, 57), (47, 55)],
+            [(3, 18), (5, 8), (6, 12), (7, 22), (15, 21), (17, 32), (19, 33), (23, 37), (26, 40), (30, 44), (31, 46), (41, 56), (42, 48), (45, 60), (51, 57), (55, 58)],
+            [(3, 16), (7, 20), (11, 26), (18, 24), (19, 25), (22, 28), (23, 29), (27, 33), (30, 36), (34, 40), (35, 41), (37, 52), (38, 44), (39, 45), (43, 56), (47, 60)],
+            [(3, 9), (7, 13), (10, 16), (11, 17), (14, 20), (15, 30), (19, 34), (21, 36), (23, 38), (25, 40), (26, 32), (27, 42), (29, 44), (31, 37), (33, 48), (43, 49), (46, 52), (47, 53), (50, 56), (54, 60)],
+            [(3, 8), (7, 10), (9, 12), (11, 18), (13, 14), (15, 24), (17, 22), (19, 28), (21, 26), (23, 25), (27, 34), (29, 36), (30, 32), (31, 33), (35, 44), (37, 42), (38, 40), (39, 48), (41, 46), (45, 52), (49, 50), (51, 54), (53, 56), (55, 60)],
+            [(3, 6), (7, 12), (11, 16), (15, 17), (18, 20), (19, 24), (21, 22), (23, 30), (25, 32), (26, 28), (27, 29), (31, 38), (33, 40), (34, 36), (35, 37), (39, 44), (41, 42), (43, 45), (46, 48), (47, 52), (51, 56), (57, 60)],
+            [(3, 5), (6, 8), (7, 9), (10, 12), (11, 13), (14, 16), (15, 18), (17, 20), (19, 21), (22, 24), (23, 26), (25, 28), (27, 30), (29, 32), (31, 34), (33, 36), (35, 38), (37, 40), (39, 41), (42, 44), (43, 46), (45, 48), (47, 49), (50, 52), (51, 53), (54, 56), (55, 57), (58, 60)],
+            [(3, 4), (7, 8), (11, 12), (13, 14), (15, 16), (17, 18), (19, 20), (21, 22), (23, 24), (25, 26), (27, 28), (29, 30), (31, 32), (33, 34), (35, 36), (37, 38), (39, 40), (41, 42), (43, 44), (45, 46), (47, 48), (49, 50), (51, 52), (55, 56), (59, 60)]
+        ],
+
+}
+
+
+@cute.jit
+def optimal_sort(
+    arr: cute.Tensor,
+    n: cutlass.Constexpr[int],
+    start: cutlass.Constexpr[int] = 0,
+    ascending: cutlass.Constexpr[bool] = True
+) -> None:
+    """
+    Optimal sorting network dispatcher.
+
+    Args:
+        arr: Array to sort
+        n: Size of array (must be power of 2 and available in networks)
+        start: Starting index (default 0)
+        ascending: Sort in ascending order (default True)
+
+    Source: https://bertdobbelaere.github.io/sorting_networks.html
+    """
+    assert n in networks
+    for level in networks[n]:
+        for i, j in level:
+            compare_and_swap(arr, start + i, start + j, ascending)

quack/sort/utils.py

@@ -0,0 +1,31 @@
+import cutlass
+import cutlass.cute as cute
+
+import quack.utils as utils
+
+
+@cute.jit
+def compare_and_swap(
+    arr: cute.Tensor, i: int, j: int, ascending: bool = True, use_selection: bool = False
+) -> None:
+    """Compare and swap elements at indices i and j in ascending or descending order."""
+    if cutlass.const_expr(use_selection):
+        a, b = arr[i], arr[j]
+        if (a > b) ^ (not ascending):
+            arr[i] = b
+            arr[j] = a
+        # if cutlass.const_expr(ascending):
+        #     if a > b:
+        #         arr[i] = b
+        #         arr[j] = a
+        # else:
+        #     if a < b:
+        #         arr[i] = b
+        #         arr[j] = a
+    else:
+        min_fn = min if cutlass.const_expr(arr.element_type != cutlass.Float32) else utils.fmin
+        max_fn = max if cutlass.const_expr(arr.element_type != cutlass.Float32) else cute.arch.fmax
+        if cutlass.const_expr(ascending):
+            arr[i], arr[j] = min_fn(arr[i], arr[j]), max_fn(arr[i], arr[j])
+        else:
+            arr[i], arr[j] = max_fn(arr[i], arr[j]), min_fn(arr[i], arr[j])