Custom flatten function using GPU acceleration?

username: drupol
cooked: <p>Hello everyone,</p>
<p>I’m currently working on a project where I need to flatten images while preserving spatial locality, and I’m exploring the use of a Hilbert curve for this purpose. I’ve been inspired by 3blue1brown’s video on the topic (<a href="https://www.youtube.com/watch?v=3s7h2MHQtxc" rel="noopener nofollow ugc">https://www.youtube.com/watch?v=3s7h2MHQtxc</a>) and would like to implement a custom flatten function that maps an image into a Hilbert curve.</p>
<p>While I understand the concept and have experimented with some basic CPU implementations, I’m now interested in accelerating this process using PyTorch and GPU support. However, I’m not entirely sure how to optimize this custom operation for GPU execution, particularly with respect to leveraging PyTorch’s capabilities.</p>
<p>Right now, the code (hereunder) I’m playing with is very slow locally but also on Google Colab.<br>
Therefore, would anyone be able to point me in the right direction on how to implement such a custom operation with efficient GPU acceleration in PyTorch? Any references, tips, or code examples would be greatly appreciated!</p>
<p>Find here the code I’m playing with:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import numpy as np
import math

# Function to compute sign of a value
def sgn(x):
    return -1 if x &lt; 0 else (1 if x &gt; 0 else 0)

# Recursive generate2d function for Gilbert curve
def generate2d(x: int, y: int, ax: int, ay: int, bx: int, by: int):
    """Recursive generation of 2D coordinates using the Gilbert space-filling curve."""

    def move_point(px: int, py: int, dx: int, dy: int):
        """Move point (px, py) in direction (dx, dy)."""
        return px + dx, py + dy

    # Width and height of the grid to fill
    w = abs(ax + ay)
    h = abs(bx + by)

    # Direction vectors
    dax, day = sgn(ax), sgn(ay)  # Major direction
    dbx, dby = sgn(bx), sgn(by)  # Orthogonal direction

    # Handle trivial row or column fills
    if h == 1 or w == 1:
        result = []
        if h == 1:
            for _ in range(w):
                result.append((x, y))
                x, y = move_point(x, y, dax, day)
        elif w == 1:
            for _ in range(h):
                result.append((x, y))
                x, y = move_point(x, y, dbx, dby)
        return result

    # Halve the movement vectors
    ax2, ay2 = ax // 2, ay // 2
    bx2, by2 = bx // 2, by // 2

    w2 = abs(ax2 + ay2)
    h2 = abs(bx2 + by2)

    result = []
    if 2 * w &gt; 3 * h:
        if w2 % 2 and w &gt; 2:
            ax2, ay2 = move_point(ax2, ay2, dax, day)

        result.extend(generate2d(x, y, ax2, ay2, bx, by))
        result.extend(generate2d(x + ax2, y + ay2, ax - ax2, ay - ay2, bx, by))

    else:
        if h2 % 2 and h &gt; 2:
            bx2, by2 = move_point(bx2, by2, dbx, dby)

        result.extend(generate2d(x, y, bx2, by2, ax2, ay2))
        result.extend(generate2d(x + bx2, y + by2, ax, ay, bx - bx2, by - by2))
        result.extend(generate2d(x + (ax - dax) + (bx2 - dbx),
                                 y + (ay - day) + (by2 - dby),
                                 -bx2, -by2, -(ax - ax2), -(ay - ay2)))

    return result

# Top-level gilbert2d function
def gilbert2d(width, height):
    if width &gt;= height:
        return generate2d(0, 0, width, 0, 0, height)
    else:
        return generate2d(0, 0, 0, height, width, 0)

# Function to reshape tensor following the Gilbert curve
def reshape_via_gilbert(tensor, width=None, height=None):
    flattened_tensor = tensor.flatten()
    num_elements = flattened_tensor.numel()

    if width is None or height is None:
        if width is None and height is not None:
            # Automatically calculate width
            width = (num_elements + height - 1) // height
        if height is None and width is not None:
            # Automatically calculate height
            height = (num_elements + width - 1) // width
        if height is None and width is None:
            # Automatically calculate width and height
            height = height or math.isqrt(num_elements)
            width = width or (num_elements + height - 1) // height

    # Create an empty tensor to store the reshaped values
    reshaped_tensor = torch.zeros((height, width), dtype=tensor.dtype, device=tensor.device)

    # Get the Gilbert curve path
    gilbert_path = gilbert2d(width, height)

    # Map the flattened tensor values to the Gilbert path
    for idx, (x, y) in enumerate(gilbert_path):
        if idx &lt; num_elements:
            reshaped_tensor[y, x] = flattened_tensor[idx]
        else:
            break

    return reshaped_tensor

# Example usage
tensor = torch.tensor([
    [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],
])

reshaped_tensor = reshape_via_gilbert(tensor)
print(reshaped_tensor)

reshaped_tensor = reshape_via_gilbert(tensor, width=5)
print(reshaped_tensor)

reshaped_tensor = reshape_via_gilbert(tensor, height=4)
print(reshaped_tensor)

reshaped_tensor = reshape_via_gilbert(tensor, width=8, height=8)
print(reshaped_tensor)
</code></pre>
<pre><code class="lang-auto">tensor([[ 1,  4,  5, 32, 33, 36],
        [ 2,  3,  6, 31, 34, 35],
        [11, 10,  7, 30, 27, 26],
        [12,  9,  8, 29, 28, 25],
        [13, 16, 17, 20, 21, 24],
        [14, 15, 18, 19, 22, 23]])
tensor([[ 1,  4,  5,  8,  9],
        [ 2,  3,  6,  7, 10],
        [19, 18, 15, 14, 11],
        [20, 17, 16, 13, 12],
        [21, 24, 25, 28, 29],
        [22, 23, 26, 27, 30],
        [ 0,  0, 35, 34, 31],
        [ 0,  0, 36, 33, 32]])
tensor([[ 1,  2, 15, 16, 17, 18, 34, 35, 36],
        [ 4,  3, 14, 13, 20, 19, 33, 32, 31],
        [ 5,  8,  9, 12, 21, 24, 25, 30, 29],
        [ 6,  7, 10, 11, 22, 23, 26, 27, 28]])
tensor([[ 1,  4,  5,  6,  0,  0,  0,  0],
        [ 2,  3,  8,  7,  0,  0,  0,  0],
        [15, 14,  9, 10,  0,  0,  0,  0],
        [16, 13, 12, 11,  0,  0,  0,  0],
        [17, 18, 31, 32, 33, 34,  0,  0],
        [20, 19, 30, 29, 36, 35,  0,  0],
        [21, 24, 25, 28,  0,  0,  0,  0],
        [22, 23, 26, 27,  0,  0,  0,  0]])
</code></pre>
<p>Generalized Hilbert algorithm credits to Jakub Červený @ <a href="https://github.com/jakubcerveny/gilbert" class="inline-onebox" rel="noopener nofollow ugc">GitHub - jakubcerveny/gilbert: Space-filling curve for rectangular domains or arbitrary size.</a></p>
<p>Thank you in advance for your help.</p>