Mojo parse failure

icp.mojo

from gpu import thread_idx, block_idx, warp, barrier
from gpu.host import DeviceContext, DeviceBuffer, HostBuffer
from gpu.memory import AddressSpace
from memory import stack_allocation
from layout import Layout, LayoutTensor
from math import sqrt
from sys import sizeof
from algorithm import parallelize

alias float32 = DType.float32
alias THREADS_PER_BLOCK = 256
alias WARP_SIZE = 32

#####
# GPU implementation of iterative closest point.
#####


fn find_nearest_neighbors_kernel[
    n: Int
](
    correspondences_tensor: LayoutTensor[
        mut=True, DType.int32, Layout.row_major(n)
    ],
    distances_tensor: LayoutTensor[mut=True, float32, Layout.row_major(n)],
    source_tensor: LayoutTensor[mut=False, float32, Layout.row_major(n, 3)],
    target_tensor: LayoutTensor[mut=False, float32, Layout.row_major(n, 3)],
):
    # Each thread processes one source point
    idx = block_idx.x * THREADS_PER_BLOCK + thread_idx.x

    # This is a guard, which prevents threads with ids
    # beyond the size of the point cloud from doing any work.
    if idx >= n:
        return

    # Load source point
    sx = source_tensor[idx, 0]
    sy = source_tensor[idx, 1]
    sz = source_tensor[idx, 2]

    var min_dist = Float32(1e10)
    var min_idx = 0

    # Find nearest neighbor in target cloud
    for j in range(n):
        tx = target_tensor[j, 0]
        ty = target_tensor[j, 1]
        tz = target_tensor[j, 2]

        dx = sx - tx
        dy = sy - ty
        dz = sz - tz
        dist_sq = dx * dx + dy * dy + dz * dz

        if dist_sq < min_dist:
            min_dist = dist_sq
            min_idx = j

    correspondences_tensor[idx] = min_idx
    distances_tensor[idx] = sqrt(min_dist)


# GPU kernel for transforming points
fn transform_points_kernel[
    n: Int
](
    transformed_tensor: LayoutTensor[mut=True, float32, Layout.row_major(3)],
    R_tensor: LayoutTensor[mut=False, float32, Layout.row_major(3, 3)],
    t_tensor: LayoutTensor[mut=False, float32, Layout.row_major(3)],
    source_tensor: LayoutTensor[mut=False, float32, Layout.row_major(3)],
):
    idx = block_idx.x * THREADS_PER_BLOCK + thread_idx.x

    # This is a guard, which prevents threads with ids
    # beyond the size of the point cloud from doing any work.
    if idx >= n:
        return

    px = source_tensor[idx, 0]
    py = source_tensor[idx, 1]
    pz = source_tensor[idx, 2]

    # Apply rotation and translation: p' = R * p + t
    x_new = (
        R_tensor[0, 0] * px
        + R_tensor[0, 1] * py
        + R_tensor[0, 2] * pz
        + t_tensor[0]
    )
    y_new = (
        R_tensor[1, 0] * px
        + R_tensor[1, 1] * py
        + R_tensor[1, 2] * pz
        + t_tensor[1]
    )
    z_new = (
        R_tensor[2, 0] * px
        + R_tensor[2, 1] * py
        + R_tensor[2, 2] * pz
        + t_tensor[2]
    )

    transformed_tensor[idx, 0] = x_new
    transformed_tensor[idx, 1] = y_new
    transformed_tensor[idx, 2] = z_new


# GPU kernel for computing centroid using warp reduction
fn compute_centroid_kernel[
    n: Int
](
    centroid_tensor: LayoutTensor[mut=True, float32, Layout.row_major(3)],
    points_tensor: LayoutTensor[mut=False, float32, Layout.row_major(n, 3)],
):
    # Each block handles one dimension of the centroid computation.
    dim = block_idx.x

    # This is a guard, which prevents threads with ids
    # beyond the size of the point cloud from doing any work.
    if dim >= 3:
        return

    # Shared stack memory for partial sums.
    var shared = stack_allocation[
        THREADS_PER_BLOCK,
        Scalar[float32],
        address_space = AddressSpace.SHARED,
    ]()

    # Each thread sums multiple elements
    tid = thread_idx.x
    idx = tid
    sum = Float32(0.0)
    while idx < n:
        sum += points_tensor[idx, dim]
        idx += THREADS_PER_BLOCK

    # Store in shared memory
    shared[tid] = sum

    # Synchronization barrier for shared memory.
    # This is required when working with shared memory!
    barrier()

    # Parallel reduction in shared memory
    var stride = THREADS_PER_BLOCK // 2
    while stride > 0:
        if tid < stride:
            shared[tid] = shared[tid] + shared[tid + stride]
        barrier()
        stride //= 2

    # Thread 0 writes final result
    if tid == 0:
        centroid_tensor[dim] = shared[0] / Float32(n)


# GPU kernel for computing cross-covariance matrix elements
fn compute_cross_covariance_kernel[
    n: Int
](
    W_tensor: LayoutTensor[mut=True, float32, Layout.row_major(3, 3)],
    correspondences_tensor: LayoutTensor[
        mut=True, DType.int32, Layout.row_major(n)
    ],
    distances_tensor: LayoutTensor[mut=True, float32, Layout.row_major(n)],
    source_centered_tensor: LayoutTensor[
        mut=False, float32, Layout.row_major(n, 3)
    ],
    target_centered_tensor: LayoutTensor[
        mut=False, float32, Layout.row_major(n, 3)
    ],
    max_dist: Float32,
):
    # Each block computes one element of the 3x3 matrix
    row = block_idx.x
    col = block_idx.y

    if row >= 3 or col >= 3:
        return

    # Shared stack memory for partial sums
    var shared = stack_allocation[
        THREADS_PER_BLOCK,
        Scalar[float32],
        address_space = AddressSpace.SHARED,
    ]()

    tid = thread_idx.x
    var sum = Float32(0)

    # Each thread processes multiple correspondences
    var idx = tid
    while idx < n:
        # Check if correspondence is valid (within distance threshold)
        if distances_tensor[idx] < max_dist:
            target_idx = correspondences_tensor[idx]
            s_val = source_centered_tensor[idx, row]
            t_val = target_centered_tensor[target_idx, col]
        idx += THREADS_PER_BLOCK

    # Store in shared memory and reduce
    shared[tid] = sum
    barrier()

    # Parallel reduction
    var stride = THREADS_PER_BLOCK // 2
    while stride > 0:
        if tid < stride:
            shared[tid] = shared[tid] + shared[tid + stride]
        barrier()
        stride //= 2

    # Thread 0 writes final result
    if tid == 0:
        W_tensor[row, col] = shared[0]


struct ICP:
    var ctx: DeviceContext
    var max_iterations: Int
    var tolerance: Float32
    var max_correspondence_dist: Float32

    def __init__(
        out self,
        max_iterations: Int = 50,
        tolerance: Float32 = 1e-6,
        max_correspondence_dist: Float32 = 1.0,
    ):
        self.ctx = DeviceContext()
        self.max_iterations = max_iterations
        self.tolerance = tolerance
        self.max_correspondence_dist = max_correspondence_dist

    fn align[
        n: Int
    ](
        self,
        source_points: HostBuffer[float32],
        target_points: HostBuffer[float32],
    ) raises -> (HostBuffer[float32], HostBuffer[float32]):
        # Create device buffers
        var source_device = self.ctx.enqueue_create_buffer[float32](n * 3)
        var target_device = self.ctx.enqueue_create_buffer[float32](n * 3)
        var transformed_device = self.ctx.enqueue_create_buffer[float32](n * 3)
        var correspondences_device = self.ctx.enqueue_create_buffer[
            DType.int32
        ](n)
        var distances_device = self.ctx.enqueue_create_buffer[float32](n)

        # Transformation matrices
        var R_device = self.ctx.enqueue_create_buffer[float32](9)
        var t_device = self.ctx.enqueue_create_buffer[float32](3)

        # Initialize R as identity and t as zero
        with R_device.map_to_host() as R_host:
            for i in range(3):
                for j in range(3):
                    R_host[i * 3 + j] = Float32(1.0) if i == j else Float32(0.0)

        _ = t_device.enqueue_fill(0)

        # Copy source and target points to device
        _ = source_device.enqueue_copy_from(source_points)
        _ = target_device.enqueue_copy_from(target_points)

        # Create tensors for easier indexing
        alias source_layout = Layout.row_major(n, 3)
        alias target_layout = Layout.row_major(n, 3)
        alias transform_layout = Layout.row_major(3, 3)
        alias vector_layout = Layout.row_major(3)

        source_tensor = LayoutTensor[float32, source_layout, MutableAnyOrigin](
            source_device
        )
        target_tensor = LayoutTensor[float32, target_layout, MutableAnyOrigin](
            target_device
        )
        transformed_tensor = LayoutTensor[
            float32, source_layout, MutableAnyOrigin
        ](transformed_device)
        correspondences_tensor = LayoutTensor[
            DType.int32, Layout.row_major(n), MutableAnyOrigin
        ](correspondences_device)
        distances_tensor = LayoutTensor[
            float32, Layout.row_major(n), MutableAnyOrigin
        ](distances_device)
        R_tensor = LayoutTensor[float32, transform_layout, MutableAnyOrigin](
            R_device
        )
        t_tensor = LayoutTensor[float32, vector_layout, MutableAnyOrigin](
            t_device
        )

        # Allocate buffers for centroids and covariance
        var source_centroid_device = self.ctx.enqueue_create_buffer[float32](3)
        var target_centroid_device = self.ctx.enqueue_create_buffer[float32](3)
        var W_device = self.ctx.enqueue_create_buffer[float32](9)

        source_centroid_tensor = LayoutTensor[
            float32, vector_layout, MutableAnyOrigin
        ](source_centroid_device)
        target_centroid_tensor = LayoutTensor[
            float32, vector_layout, MutableAnyOrigin
        ](target_centroid_device)
        W_tensor = LayoutTensor[float32, transform_layout, MutableAnyOrigin](
            W_device
        )

        # ICP iterations
        var prev_error = Float32(1e10)
        grid_dim_points = (n + THREADS_PER_BLOCK - 1) // THREADS_PER_BLOCK

        for iteration in range(self.max_iterations):
            # Transform source points
            self.ctx.enqueue_function[transform_points_kernel](
                source_tensor,
                transformed_tensor,
                R_tensor,
                t_tensor,
                n,
                grid_dim=grid_dim_points,
                block_dim=THREADS_PER_BLOCK,
            )

            # Find nearest neighbors
            self.ctx.enqueue_function[find_nearest_neighbors_kernel](
                transformed_tensor,
                target_tensor,
                correspondences_tensor,
                distances_tensor,
                n,
                n,
                grid_dim=grid_dim_points,
                block_dim=THREADS_PER_BLOCK,
            )

            # Compute centroids
            self.ctx.enqueue_function[compute_centroid_kernel](
                transformed_tensor,
                source_centroid_tensor,
                n,
                grid_dim=3,
                block_dim=THREADS_PER_BLOCK,
            )

            # For target centroid, we need a custom kernel that considers correspondences
            # (simplified here - in practice you'd write a specialized kernel)
            self.ctx.enqueue_function[compute_centroid_kernel](
                target_tensor,
                target_centroid_tensor,
                n,
                grid_dim=3,
                block_dim=THREADS_PER_BLOCK,
            )

            # Center the point clouds (would need centering kernels)
            # Compute cross-covariance matrix
            self.ctx.enqueue_function[compute_cross_covariance_kernel](
                transformed_tensor,
                target_tensor,
                W_tensor,
                correspondences_tensor,
                distances_tensor,
                self.max_correspondence_dist,
                n,
                grid_dim=(3, 3),
                block_dim=THREADS_PER_BLOCK,
            )

            # SVD and rotation update would happen here
            # For now, this is a simplified version

            # Check convergence (would need error computation kernel)
            self.ctx.synchronize()

            # Simple convergence check - in practice you'd compute actual error
            if iteration > 5:  # Simplified termination
                break

        # Copy final transformation back to host
        var R_host = HostBuffer[float32](self.ctx, 9)
        var t_host = HostBuffer[float32](self.ctx, 3)

        with R_device.map_to_host() as R_map:
            for i in range(9):
                R_host[i] = R_map[i]

        with t_device.map_to_host() as t_map:
            for i in range(3):
                t_host[i] = t_map[i]

        return (R_host, t_host)


# Example usage
def main():
    # Create example point clouds
    alias n_points = 1000
    var ctx = DeviceContext()
    var source_points = HostBuffer[float32](ctx, 9)
    var target_points = HostBuffer[float32](ctx, 9)

    # Initialize with some example data
    for i in range(n_points):
        # Source points
        source_points[i * 3] = Float32(i) / Float32(n_points)
        source_points[i * 3 + 1] = Float32(i * 2) / Float32(n_points)
        source_points[i * 3 + 2] = Float32(i * 3) / Float32(n_points)

        # Target points (slightly transformed)
        target_points[i * 3] = source_points[i * 3] + 0.1
        target_points[i * 3 + 1] = source_points[i * 3 + 1] + 0.2
        target_points[i * 3 + 2] = source_points[i * 3 + 2] + 0.3

    # Run ICP
    icp = ICP()
    (R, t) = icp.align[n_points](source_points, target_points)

    print("Rotation matrix:")
    for i in range(3):
        print(R[i * 3], R[i * 3 + 1], R[i * 3 + 2])

    print("Translation vector:")
    print(t[0], t[1], t[2])