Torch Compile and External Kernels

            

            
class MLP(torch.nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        self.fc1 = torch.nn.Linear(input_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = torch.nn.Linear(hidden_dim, output_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
A simple 3 layer MLP that takes in a tensor of
shape (N, input_dim) and outputs a tensor of
shape (N, output_dim)
"""
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x
        

            

            
class kNN_Projector(torch.nn.Module):
    def __init__(self, k: int, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        self.proj = MLP(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=hidden_dim)
        self.k = k

        self.proj_out = MLP(input_dim=hidden_dim, hidden_dim = hidden_dim, output_dim = output_dim)

    def forward(self, p1: torch.Tensor, p2: torch.Tensor) -> torch.Tensor:
        """
Accept two point clouds, p1 and p2.  Compute a learnable projection onto p2 to
learn features.  Then, use a kNN-weighted aggregation to project those features
onto p1.
"""
        p2_features = self.proj(p2)

        p1_features = knn_weighted_feature_aggregation(p1, p2, p2_features, k=self.k)

        return self.proj_out(p1_features)
        

            

            
def knn_weighted_feature_aggregation(
    p1: torch.Tensor,
    p2: torch.Tensor,
    p2_features: torch.Tensor,
    k: int = 3,
    sigma: float = 0.1,
    eps: float = 1e-8
    ) -> torch.Tensor:
    """
Perform differentiable kNN-weighted feature aggregation.

Args:
p1 (torch.Tensor): Query points, shape (B, M, D)
p2 (torch.Tensor): Reference points, shape (B, N, D)
p2_features (torch.Tensor): Features at reference points, shape (B, N, D_feat)
k (int): Number of neighbors
sigma (float): RBF temperature parameter
eps (float): Numerical stability for normalization

Returns:
torch.Tensor: Aggregated features at p1, shape (B, M, D_feat)
"""
    # M, D = p1.shape
    # N, D_feat = p2_features.shape

    # Compute pairwise distances: (M, N)
    dists = torch.norm(p1[:,None,:] - p2[None,:,:], dim=-1)


    # Find top-k nearest neighbors
    topk_dists, topk_idx = torch.topk(dists, k=k, dim=1, largest=False)

    # Gather neighbor features: (M, k, D_feat)
    neighbors = p2_features[topk_idx]

    # Compute weights: (M, k)
    weights = torch.softmax(-topk_dists / sigma, dim=1)

    # Weighted sum of neighbor features: (M, D_feat)
    agg = torch.sum(weights.unsqueeze(-1) * neighbors, dim=1)

    return agg
        

            

            
def generate_data(N_points_to_search, grid_points, target_features, dtype=torch.bfloat16):
    device = torch.device("cuda")


    # Make a random point cloud:
    point_cloud = torch.randn(N_points_to_search, 3, device=device, requires_grad=False, dtype=dtype)

    # And this is a set of 3D points on a grid, that we'll flatten:
    x = torch.linspace(-1, 1, 30, device=device, dtype=dtype)
    y = torch.linspace(-1, 1, 30, device=device, dtype=dtype)
    z = torch.linspace(-1, 1, 30, device=device, dtype=dtype)

    # Create 3D meshgrid
    X, Y, Z = torch.meshgrid(x, y, z, indexing='ij')

    # Flatten and stack to get grid points as (N, 3) tensor
    grid_points = torch.stack([X.flatten(), Y.flatten(), Z.flatten()], dim=1)

    grid_features = torch.randn(grid_points.shape[0], target_features, device=device, requires_grad=False, dtype=dtype)

    return point_cloud, grid_points, grid_features

def set_seed(seed: int):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def train_step(model, optimizer, grid_points, point_cloud, grid_features):
    # Pretent to train the model!
    optimizer.zero_grad()
    output = model.forward(grid_points, point_cloud)
    loss = torch.mean((output - grid_features)**2)
    loss.backward()
    optimizer.step()
    return loss
        

            

            
def measure_performance(model, inputs, warmup_iters, benchmark_iters, profile=False):

    grid_points, point_cloud, grid_features = inputs

    profiler = Profiler()
    if profile:
        profiler.enable("torch")

    # Make a dummy optimizer:
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # Warm up:
    for i in range(warmup_iters):
        # Forward only:
        model.forward(grid_points, point_cloud)

        # Training:
        loss = train_step(model, optimizer, grid_points, point_cloud, grid_features)


    torch.cuda.synchronize()

    with profiler:


        with torch.no_grad():
            torch.cuda.synchronize()
            start_event = torch.cuda.Event(enable_timing=True)
            end_event = torch.cuda.Event(enable_timing=True)

            start_event.record()
            # Benchmark the forward pass
            for i in range(benchmark_iters):
                output =model.forward(grid_points, point_cloud)
            end_event.record()
            torch.cuda.synchronize()

        print(f"Time taken in forward:{start_event.elapsed_time(end_event) / benchmark_iters:.3f}ms per iteration")

        # Benchmark the training loop:

        torch.cuda.synchronize()
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)

        start_event.record()
        # Benchmark the backward pass
        for i in range(benchmark_iters):
            loss = train_step(model, optimizer, grid_points, point_cloud, grid_features)

        end_event.record()
        torch.cuda.synchronize()


        print(f"Time taken in backward:{start_event.elapsed_time(end_event) / benchmark_iters:.3f}ms per iteration")
        

            

            
if __name__ == "__main__":

    set_seed(42)

    target_features = 1
    n_grid_points = 30
    n_cloud_points = 100000
    dtype = torch.float32

    point_cloud, grid_points, grid_features = generate_data(n_cloud_points, n_grid_points, target_features, dtype=dtype)
    print(point_cloud.shape)
    print(grid_points.shape)
    print(grid_features.shape)

    model = kNN_Projector(k=7, hidden_dim=25, output_dim=target_features).cuda().to(dtype)

    warmup_iters = 5
    benchmark_iters = 15

    measure_performance(model, (grid_points, point_cloud, grid_features), warmup_iters, benchmark_iters, profile=False)
        

            

            
torch.Size([100000, 3])
torch.Size([27000, 3])
torch.Size([27000, 1])
Time taken in forward: 144.045 ms per iteration
Time taken in backward: 144.758 ms per iteration
        

            

            
Time taken in forward: 74.237 ms per iteration
Time taken in backward: 74.657 ms per iteration
        

            

            
def knn_weighted_feature_aggregation(
    p1: torch.Tensor,
    p2: torch.Tensor,
    p2_features: torch.Tensor,
    k: int = 3,
    sigma: float = 0.1,
    eps: float = 1e-8
    ) -> torch.Tensor:
    """
"""


    # Find top-k nearest neighbors (Make sure to cast to float32 for cuml)
    topk_dists, topk_idx = knn_search_with_cuml(p1.to(torch.float32), p2.to(torch.float32), k)

    # Gather neighbor features: (M, k, D_feat)
    neighbors = p2_features[topk_idx]

    # Compute weights: (M, k)
    weights = torch.softmax(-topk_dists / sigma, dim=1)

    # Weighted sum of neighbor features: (M, D_feat)
    agg = torch.sum(weights.unsqueeze(-1) * neighbors, dim=1)

    # Cast back to original dtype
    return agg.to(p1.dtype)
        

            

            
def knn_search_with_cuml(p1: torch.Tensor, p2: torch.Tensor, k: int = 3):
    # Use dlpack to move the data without copying between pytorch and cuml:
    p1 = cp.from_dlpack(p1)
    p2 = cp.from_dlpack(p2)

    # Construct the knn:
    knn = cuml.neighbors.NearestNeighbors(n_neighbors=k)
    # First pass partitions everything in p2 to make lookups fast
    knn.fit(p2)

    # Second pass uses that partition to quickly find neighbors of points in p1
    distance, indices = knn.kneighbors(p1)

    # convert back to pytorch:
    distance = torch.from_dlpack(distance)
    indices = torch.from_dlpack(indices)

    # Return torch objects.
    return distance, indices
        

            

            
Time taken in forward: 14.139 ms per iteration
Time taken in backward: 15.842 ms per iteration
        

            

            
/usr/local/lib/python3.12/dist-packages/torch/_dynamo/variables/functions.py:700: UserWarning: Graph break due to unsupported builtin cupy._core.dlpack.from_dlpack. This function is either a Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python extension (perhaps created with pybind). If it is a Python builtin, please file an issue on GitHub so the PyTorch team can add support for it and see the next case for a workaround. If it is a third-party C/C++ Python extension, please either wrap it into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html for more details) or, if it is traceable, use torch.compiler.allow_in_graph.
        

            

            
Time taken in forward: 15.697 ms per iteration
Time taken in backward: 17.670 ms per iteration
        

            

            
@torch.library.custom_op("cuml::knn", mutates_args=())
def knn_search_with_cuml(p1: torch.Tensor, p2: torch.Tensor, k: int = 3) -> tuple[torch.Tensor, torch.Tensor]:
    p1 = cp.from_dlpack(p1)
    p2 = cp.from_dlpack(p2)

    knn = cuml.neighbors.NearestNeighbors(n_neighbors=k)
    knn.fit(p2)

    distance, indices = knn.kneighbors(p1)

    # convert back to pytorch:
    distance = torch.from_dlpack(distance)
    indices = torch.from_dlpack(indices)

    return distance, indices
        

            

            
@knn_search_with_cuml.register_fake
def _(p1, p2, k):
    assert p1.device == p2.device

    dist_output = torch.empty(p1.shape[0], k, device=p1.device, dtype=p1.dtype)
    idx_output = torch.empty(p1.shape[0], k, device=p1.device, dtype=torch.int64)

    return dist_output, idx_output
        

            

            
import rmm
from rmm.allocators.torch import rmm_torch_allocator
        

            

            
rmm.reinitialize(pool_allocator=True)
torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
        

            

            
# Not Compiled:
Time taken in forward: 11.185 ms per iteration
Time taken in backward: 12.818 ms per iteration

# Compiled:
Time taken in forward: 11.163 ms per iteration
Time taken in backward: 12.571 ms per iteration
        

            

            
def radius_bounded_feature_aggregation(
    p1: torch.Tensor,
    p2: torch.Tensor,
    p2_features: torch.Tensor,
    radius: float,
    sigma: float,
    ) -> torch.Tensor:
    """
Perform differentiable radius-bounded feature aggregation.

Args:
p1 (torch.Tensor): Query points, shape (B, M, D)
p2 (torch.Tensor): Reference points, shape (B, N, D)
p2_features (torch.Tensor): Features at reference points, shape (B, N, D_feat)
radius (float): Radius for neighbor search
sigma (float): RBF temperature parameter

Returns:
torch.Tensor: Aggregated features at p1, shape (B, M, D_feat)
"""

    # Compute pairwise distances: (M, N)
    dists = torch.norm(p1[:,None,:] - p2[None,:,:], dim=-1)

    # Create mask for neighbors within radius
    mask = dists <= radius

    # Compute weights from all distances first
    weights = torch.softmax(-dists / sigma, dim=-1)

    # Apply mask to zero out weights for points outside radius
    weights = torch.where(mask, weights, torch.zeros_like(weights))

    # Renormalize weights so they sum to 1 for each query point
    weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8)

    # Weighted sum of all reference features: (M, D_feat)
    agg = torch.sum(weights.unsqueeze(-1) * p2_features.unsqueeze(0), dim=1)

    return agg
        

            

            
Time taken in forward: 161.396 ms per iteration
Time taken in backward: 260.719 ms per iteration
        

            

            
Time taken in forward: 38.718 ms per iteration
Time taken in backward: 86.226 ms per iteration
        

            

            
@wp.kernel
def warp_radius_bounded_feature_aggregation(
    query_points: wp.array(dtype=wp.vec3),
    ref_points:  wp.array(dtype=wp.vec3),
    ref_features: wp.array2d(dtype=wp.float32),
    grid: wp.uint64,
    radius: float,
    sigma: float,
    output_features: wp.array2d(dtype=wp.float32),
):
    # Get the thread ID:
    tid = wp.tid()

    # Get position from query points
    pos = query_points[tid]

    feature_dim = ref_features.shape[1]
    local_output = output_features[tid]

    # Find all the neighbors using the hash grid:
    neighbors = wp.hash_grid_query(id=grid, point=pos, max_dist=radius)

    weight = float(0.0)
    # Loop through neighbors. Compute a weighted distance and accumulate it,
    # but also track the weight to normalize at the end.
    for index in neighbors:
        # Get the neighbor position:
        pos2 = ref_points[index]

        # Compute the distance:

        dist = wp.length(pos - pos2)
        # disp = pos - pos2
        # dist = wp.sqrt(disp[0]**2. + disp[1]**2. + disp[2]**2.)

        if dist > radius:
            continue

        # Get the features at this index:
        feature = ref_features[index]

        # Compute the weight:
        this_weight = wp.exp( - dist / sigma)

        # Accumulate the weight, and weight * feature
        weight += this_weight
        # Work directly with the 2D array indexing instead of the slice
        for j in range(feature_dim):
            local_output[j] += feature[j] * this_weight

    if weight > 0.0:
        # Normalize by working directly with the 2D array
        for j in range(feature_dim):
            local_output[j] /= (weight + 1e-8)

    # Write the output
    for j in range(feature_dim):
        output_features[tid,j] = local_output[j]
        

            

            
@torch.library.custom_op("warp::radius_bounded_feature_aggregation", mutates_args=())
def radius_bounded_feature_aggregation_impl(
    query_points: torch.Tensor,
    ref_points: torch.Tensor,
    ref_features: torch.Tensor,
    radius: float,
    sigma: float,
) -> torch.Tensor:


    # Convert to warp
    # We can only build and query the points in wp.vec3 format:
    wp_query_points = wp.from_torch(
        query_points,
        dtype=wp.vec3, # vec3 here!
        requires_grad=query_points.requires_grad
    )
    wp_ref_points = wp.from_torch(
        ref_points,
        dtype=wp.vec3, # and here!
        requires_grad=ref_points.requires_grad
    )
    wp_ref_features = wp.from_torch(
        ref_features,
        dtype=wp.float32, # but, the features we keep as an array!
        requires_grad=ref_features.requires_grad
    )


    # In the data generation, we had set the grid to range over -1 to 1 with 30 points.
    # We can use that to dictate the grid size:
    grid_size = int(2. / radius) + 1

    # **In general** to do this, you'd have to incur a performance penalty:
    # 1. Find the min max of all the points in your query set.
    # 2. Divide the range by the radius to get the grid size.
    # 3. Move the grid size to the CPU (which is blocking if done wrong)
    # 4. Construct the grid on the CPU.

    # But, it's **likely** you'd be constructing the grid once and caching it, anyways.



    # Build the grid used in the kernel:
    hash_grid = wp.HashGrid(grid_size, grid_size, grid_size)

    # This actually loops over the points and does the hashing
    hash_grid.build(wp_ref_points, radius)


    # Allocate output space (with pytorch!) and convert to warp:
    output_features = torch.zeros(
        (query_points.shape[0], ref_features.shape[1], ),
        device=ref_features.device,
        dtype=ref_features.dtype,
        requires_grad=ref_features.requires_grad
    )

    wp_output = wp.from_torch(
        output_features,
        dtype=wp.float32,
        requires_grad=ref_features.requires_grad
    )

    # Launch the kernel:
    feature_dim = ref_features.shape[1]
    wp.launch(
        warp_radius_bounded_feature_aggregation,
        inputs=[
            wp_query_points,
            wp_ref_points,
            wp_ref_features,
            hash_grid.id,
            radius,
            sigma,
        ],
        outputs =[
            wp_output,
        ],
        dim=[query_points.shape[0]],
    )


    # return the output features:
    return output_features
        

            

            
@radius_bounded_feature_aggregation_impl.register_fake
def _(
    query_points: torch.Tensor,
    ref_points: torch.Tensor,
    ref_features: torch.Tensor,
    radius: float,
    sigma: float,
) -> torch.Tensor:

    assert query_points.is_cuda
    assert ref_points.is_cuda
    assert ref_features.is_cuda

    output = torch.empty(
        (query_points.shape[0], ref_features.shape[1], ),
        device=query_points.device,
        dtype=query_points.dtype)

    return output
        

            

            
@torch.library.custom_op("warp::radius_bounded_feature_aggregation_bwd", mutates_args=())
def radius_bounded_feature_aggregation_bwd_impl(
    query_points: torch.Tensor,
    ref_points: torch.Tensor,
    ref_features: torch.Tensor,
    radius: float,
    sigma: float,
    output_features: torch.Tensor,
    grad_outputs : torch.Tensor,
) -> torch.Tensor:

    # This function only needs to get the gradients of p2_features,
    # based on the grad_output of the forward outputs.
    # Everything else is None for a grad.

    # Convert to warp:
    wp_query_points = wp.from_torch(
        query_points,
        dtype=wp.vec3,
        requires_grad=False
    )
    # We can only build and query the points in float32:
    wp_ref_points = wp.from_torch(
        ref_points,
        dtype=wp.vec3,
        requires_grad=False
    )
    #########################################################
    # Because we set requires_grad True here, warp will
    # populate gradients HERE in the .grad attribute.
    #########################################################
    wp_ref_features = wp.from_torch(
        ref_features,
        dtype=wp.float32,
        requires_grad=True,
    )

    # In the data generation below, we set the grid to range over -1 to 1 with 30 points.
    # We can use that to dictate the grid size:
    grid_size = int(2. / radius) + 1

    # Build the grid used in the kernel:
    # In a real application, you'd cache and retrieve this in the backwards pass.
    # The trick would be to make sure the hash_grid object persists but is not
    # actually in the torch interface (it would break things for the compiler!)
    hash_grid = wp.HashGrid(grid_size, grid_size, grid_size)

    # We're rebuilding here just to make the implementation straightforward
    hash_grid.build(wp_ref_points, radius)

    wp_output = wp.from_torch(output_features)

    wp_grad_outputs = wp.from_torch(grad_outputs)


    feature_dim = ref_features.shape[1]
    # Launch the kernel:
    wp.launch(
        warp_radius_bounded_feature_aggregation,
        inputs=[
            wp_query_points,
            wp_ref_points,
            wp_ref_features,
            hash_grid.id,
            radius,
            sigma,
        ],
        outputs =[
            wp_output,
        ],
        adj_inputs = [
            None,
            None,
            wp_ref_features.grad,
            None,
            None,
            None,
        ],
        adj_outputs = [
            wp_grad_outputs,
        ],
        adjoint=True,  ############ Pay attention here! Launch the kernel adjoint!
        dim=[query_points.shape[0]],
    )

    # return the gradient of the features:
    return ref_features.grad
        

            

            
@radius_bounded_feature_aggregation_bwd_impl.register_fake
def _(
    query_points: torch.Tensor,
    ref_points: torch.Tensor,
    ref_features: torch.Tensor,
    radius: float,
    sigma: float,
    output_features: torch.Tensor,
    grad_outputs : torch.Tensor
) -> torch.Tensor:
    grad_outputs = torch.empty_like(ref_features)
    return grad_outputs
        

            

            
# Use this to save any inputs or outputs for the backward pass.
def setup_context(ctx: torch.autograd.function.FunctionCtx, inputs, output):
    query_points, ref_points, ref_features, radius, sigma = inputs
    out_features = output
    ctx.save_for_backward(query_points, ref_points, ref_features, out_features)

    ctx.radius = radius
    ctx.sigma = sigma

# And this connects the forward pass to the backward pass with pytorch.
radius_bounded_feature_aggregation_impl.register_autograd(radius_bounded_feature_aggregation_backward_worker, setup_context=setup_context)
        

            

            
Time taken in forward: 13.511 ms per iteration
Time taken in backward: 29.953 ms per iteration