reasoning
reasoning (truncated 13004 chars): Let me analyze this problem:
1. We need to implement a batched SAT (Separating Axis Test) for collision detection between an axis-aligned rectangle (AABB) and a triangle in 2D.
2. The reference implementation:
- Takes robot_min (B, 2), robot_max (B, 2), env_triangle (B, 3, 2)
- Computes 4 corners of the rectangle
- Tests 5 axes: 2 world axes (x, y) + 3 triangle edge normals
- For each axis, projects both shapes and checks for separation
- Returns 1.0 if overlapping, 0.0 if separated
3. Key operations:
- Computing rectangle corners from min/max
- Projecting points onto axes (dot product)
- Finding min/max of projections
- Checking separation conditions
- Normalizing edge normals
Let me write a CUDA kernel that:
1. Each thread handles one batch element
2. Computes the 4 corners locally
3. Tests all 5 axes
4. Returns the collision result
Here's my implementation plan:
```cpp
__global__ void sat_rect_triangle_kernel(
const float* robot_min, // (B, 2)
const float* robot_max, // (B, 2)
const float* env_triangle, // (B, 3, 2)
float* output, // (B, 1)
float eps,
int B
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= B) return;
// Load rectangle bounds
float x0 = robot_min[idx * 2 + 0];
float y0 = robot_min[idx * 2 + 1];
float x1 = robot_max[idx * 2 + 0];
float y1 = robot_max[idx * 2 + 1];
// Rectangle corners
float corners[4][2] = {{x0, y0}, {x1, y0}, {x1, y1}, {x0, y1}};
// Load triangle vertices
float tri[3][2];
for (int i = 0; i < 3; i++) {
tri[i][0] = env_triangle[idx * 6 + i * 2 + 0];
tri[i][1] = env_triangle[idx * 6 + i * 2 + 1];
}
bool hit = true;
// Test world axes (x and y)
// X axis
float r_min = min(x0, x1);
float r_max = max(x0, x1);
float t_min = min(min(tri[0][0], tri[1][0]), tri[2][0]);
float t_max = max(max(tri[0][0], tri[1][0]), tri[2][0]);
if (r_max < t_min - eps || t_max < r_min - eps) hit = false;
// Y axis
r_min = min(y0, y1);
r_max = max(y0, y1);
t_min = min(min(tri[0][1], tri[1][1]), tri[2][1]);
t_max = max(max(tri[0][1], tri[1][1]), tri[2][1]);
if (r_max < t_min - eps || t_max < r_min - eps) hit = false;
// Test triangle edge normals
for (int i = 0; i < 3 && hit; i++) {
// Edge from vertex i to vertex (i+1)%3
int j = (i + 1) % 3;
float ex = tri[j][0] - tri[i][0];
float ey = tri[j][1] - tri[i][1];
// Normal (perpendicular)
float nx = -ey;
float ny = ex;
// Normalize
float len = sqrtf(nx * nx + ny * ny) + 1e-8f;
nx /= len;
ny /= len;
// Project rectangle corners
float r0 = corners[0][0] * nx + corners[0][1] * ny;
float r1 = corners[1][0] * nx + corners[1][1] * ny;
float r2 = corners[2][0] * nx + corners[2][1] * ny;
float r3 = corners[3][0] * nx + corners[3][1] * ny;
float r_proj_min = min(min(r0, r1), min(r2, r3));
float r_proj_max = max(max(r0, r1), max(r2, r3));
// Project triangle vertices
float t0 = tri[0][0] * nx + tri[0][1] * ny;
float t1 = tri[1][0] * nx + tri[1][1] * ny;
float t2 = tri[2][0] * nx + tri[2][1] * ny;
float t_proj_min = min(min(t0, t1), t2);
float t_proj_max = max(max(t0, t1), t2);
if (r_proj_max < t_proj_min - eps || t_proj_max < r_proj_min - eps) hit = false;
}
output[idx] = hit ? 1.0f : 0.0f;
}
```
Actually, let me think about the optimization more carefully. Since x0 <= x1 and y0 <= y1 by definition (min and max), I can simplify the world axis tests.
Also, I need to be careful with the tensor strides - PyTorch tensors are typically row-major.
Let me write a cleaner implementation:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void sat_rect_triangle_kernel(
const float* __restrict__ robot_min,
const float* __restrict__ robot_max,
const float* __restrict__ env_triangle,
float* __restrict__ output,
float eps,
int B
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= B) return;
// Load rectangle bounds
float x0 = robot_min[idx * 2];
float y0 = robot_min[idx * 2 + 1];
float x1 = robot_max[idx * 2];
float y1 = robot_max[idx * 2 + 1];
// Load triangle vertices (B, 3, 2)
float t0x = env_triangle[idx * 6];
float t0y = env_triangle[idx * 6 + 1];
float t1x = env_triangle[idx * 6 + 2];
float t1y = env_triangle[idx * 6 + 3];
float t2x = env_triangle[idx * 6 + 4];
float t2y = env_triangle[idx * 6 + 5];
bool hit = true;
// Test X axis
float r_x_min = x0, r_x_max = x1;
float t_x_min = fminf(fminf(t0x, t1x), t2x);
float t_x_max = fmaxf(fmaxf(t0x, t1x), t2x);
if (r_x_max < t_x_min - eps || t_x_max < r_x_min - eps) hit = false;
// Test Y axis
if (hit) {
float r_y_min = y0, r_y_max = y1;
float t_y_min = fminf(fminf(t0y, t1y), t2y);
float t_y_max = fmaxf(fmaxf(t0y, t1y), t2y);
if (r_y_max < t_y_min - eps || t_y_max < r_y_min - eps) hit = false;
}
// Test triangle edge normals
for (int i = 0; i < 3 && hit; i++) {
float ex, ey;
if (i == 0) {
ex = t1x - t0x;
ey = t1y - t0y;
} else if (i == 1) {
ex = t2x - t1x;
ey = t2y - t1y;
} else {
ex = t0x - t2x;
ey = t0y - t2y;
}
// Normal (perpendicular)
float nx = -ey;
float ny = ex;
// Normalize
float len_sq = nx * nx + ny * ny;
float inv_len = rsqrtf(len_sq + 1e-8f);
nx *= inv_len;
ny *= inv_len;
// Project rectangle corners
float c0 = x0 * nx + y0 * ny;
float c1 = x1 * nx + y0 * ny;
float c2 = x1 * nx + y1 * ny;
float c3 = x0 * nx + y1 * ny;
float r_proj_min = fminf(fminf(c0, c1), fminf(c2, c3));
float r_proj_max = fmaxf(fmaxf(c0, c1), fmaxf(c2, c3));
// Project triangle vertices
float p0 = t0x * nx + t0y * ny;
float p1 = t1x * nx + t1y * ny;
float p2 = t2x * nx + t2y * ny;
float t_proj_min = fminf(fminf(p0, p1), p2);
float t_proj_max = fmaxf(fmaxf(p0, p1), p2);
if (r_proj_max < t_proj_min - eps || t_proj_max < r_proj_min - eps) hit = false;
}
output[idx] = hit ? 1.0f : 0.0f;
}
torch::Tensor sat_rect_triangle_cuda(
torch::Tensor robot_min,
torch::Tensor robot_max,
torch::Tensor env_triangle,
float eps
) {
int B = robot_min.size(0);
auto output = torch::empty({B, 1}, robot_min.options());
int threads = 256;
int blocks = (B + threads - 1) / threads;
sat_rect_triangle_kernel<<<blocks, threads>>>(
robot_min.data_ptr<float>(),
robot_max.data_ptr<float>(),
env_triangle.data_ptr<float>(),
output.data_ptr<float>(),
eps,
B
);
return output;
}
"""
sat_module = load_inline(
name="sat_rect_triangle",
cpp_sources="",
cuda_sources=cuda_source,
functions=["sat_rect_triangle_cuda"],
extra_cuda_cflags=["-O3"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, eps: float = 1e-6):
super().__init__()
self.eps = eps
def forward(
self,
robot_min: torch.Tensor,
robot_max: torch.Tensor,
env_triangle: torch.Tensor,
) -> torch.Tensor:
return sat_module.sat_rect_triangle_cuda(robot_min, robot_max, env_triangle, self.eps)
```
Wait, I need to double-check the reference implementation logic more carefully.
Looking at the reference:
1.
...[ omitted 13004 characters ]...
min, t_max;
project_rect(x0, y0, x1, y1, nx, ny, r_min, r_max);
project_tri(t0x, t0y, t1x, t1y, t2x, t2y, nx, ny, t_min, t_max);
if (r_max < t_min - eps || t_max < r_min - eps) hit = false;
}
output[idx] = hit ? 1.0f : 0.0f;
}
torch::Tensor sat_rect_triangle_cuda(
torch::Tensor robot_min,
torch::Tensor robot_max,
torch::Tensor env_triangle,
float eps
) {
const int B = robot_min.size(0);
auto output = torch::empty({B, 1}, robot_min.options());
const int threads = 256;
const int blocks = (B + threads - 1) / threads;
sat_rect_triangle_kernel<<<blocks, threads>>>(
robot_min.data_ptr<float>(),
robot_max.data_ptr<float>(),
env_triangle.data_ptr<float>(),
output.data_ptr<float>(),
eps,
B
);
return output;
}
"""
sat_module = load_inline(
name="sat_rect_triangle_v3",
cpp_sources="",
cuda_sources=cuda_source,
functions=["sat_rect_triangle_cuda"],
extra_cuda_cflags=["-O3", "-use_fast_math"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, eps: float = 1e-6):
super().__init__()
self.eps = eps
def forward(
self,
robot_min: torch.Tensor,
robot_max: torch.Tensor,
env_triangle: torch.Tensor,
) -> torch.Tensor:
robot_min = robot_min.contiguous()
robot_max = robot_max.contiguous()
env_triangle = env_triangle.contiguous()
return sat_module.sat_rect_triangle_cuda(robot_min, robot_max, env_triangle, self.eps)
```
Actually, I realize using a 2D array `edges[3][2]` in device code might not be ideal for register allocation. Let me just compute the edges on the fly within the loop, which is cleaner and likely better for the compiler:
Actually, let me simplify and just do direct computation. The device functions might also introduce unnecessary overhead. Let me flatten everything into one clean kernel:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void sat_rect_triangle_kernel(
const float* __restrict__ robot_min,
const float* __restrict__ robot_max,
const float* __restrict__ env_triangle,
float* __restrict__ output,
float eps,
int B
) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= B) return;
// Load rectangle bounds
const float x0 = robot_min[idx * 2];
const float y0 = robot_min[idx * 2 + 1];
const float x1 = robot_max[idx * 2];
const float y1 = robot_max[idx * 2 + 1];
// Load triangle vertices
const float t0x = env_triangle[idx * 6];
const float t0y = env_triangle[idx * 6 + 1];
const float t1x = env_triangle[idx * 6 + 2];
const float t1y = env_triangle[idx * 6 + 3];
const float t2x = env_triangle[idx * 6 + 4];
const float t2y = env_triangle[idx * 6 + 5];
bool hit = true;
// Test X axis [1, 0]
const float t_x_min = fminf(fminf(t0x, t1x), t2x);
const float t_x_max = fmaxf(fmaxf(t0x, t1x), t2x);
if (x1 < t_x_min - eps || t_x_max < x0 - eps) {
output[idx] = 0.0f;
return;
}
// Test Y axis [0, 1]
const float t_y_min = fminf(fminf(t0y, t1y), t2y);
const float t_y_max = fmaxf(fmaxf(t0y, t1y), t2y);
if (y1 < t_y_min - eps || t_y_max < y0 - eps) {
output[idx] = 0.0f;
return;
}
// Edge 0: v0 -> v1, normal = perpendicular (rotated 90 CCW)
{
const float ex = t1x - t0x;
const float ey = t1y - t0y;
const float len = sqrtf(ex * ex + ey * ey) + 1e-8f;
const float nx = -ey / len;
const float ny = ex / len;
const float c0 = x0 * nx + y0 * ny;
const float c1 = x1 * nx + y0 * ny;
const float c2 = x1 * nx + y1 * ny;
const float c3 = x0 * nx + y1 * ny;
const float r_min = fminf(fminf(c0, c1), fminf(c2, c3));
const float r_max = fmaxf(fmaxf(c0, c1), fmaxf(c2, c3));
const float p0 = t0x * nx + t0y * ny;
const float p1 = t1x * nx + t1y * ny;
const float p2 = t2x * nx + t2y * ny;
const float t_min = fminf(fminf(p0, p1), p2);
const float t_max = fmaxf(fmaxf(p0, p1), p2);
if (r_max < t_min - eps || t_max < r_min - eps) {
output[idx] = 0.0f;
return;
}
}
// Edge 1: v1 -> v2
{
const float ex = t2x - t1x;
const float ey = t2y - t1y;
const float len = sqrtf(ex * ex + ey * ey) + 1e-8f;
const float nx = -ey / len;
const float ny = ex / len;
const float c0 = x0 * nx + y0 * ny;
const float c1 = x1 * nx + y0 * ny;
const float c2 = x1 * nx + y1 * ny;
const float c3 = x0 * nx + y1 * ny;
const float r_min = fminf(fminf(c0, c1), fminf(c2, c3));
const float r_max = fmaxf(fmaxf(c0, c1), fmaxf(c2, c3));
const float p0 = t0x * nx + t0y * ny;
const float p1 = t1x * nx + t1y * ny;
const float p2 = t2x * nx + t2y * ny;
const float t_min = fminf(fminf(p0, p1), p2);
const float t_max = fmaxf(fmaxf(p0, p1), p2);
if (r_max < t_min - eps || t_max < r_min - eps) {
output[idx] = 0.0f;
return;
}
}
// Edge 2: v2 -> v0
{
const float ex = t0x - t2x;
const float ey = t0y - t2y;
const float len = sqrtf(ex * ex + ey * ey) + 1e-8f;
const float nx = -ey / len;
const float ny = ex / len;
const float c0 = x0 * nx + y0 * ny;
const float c1 = x1 * nx + y0 * ny;
const float c2 = x1 * nx + y1 * ny;
const float c3 = x0 * nx + y1 * ny;
const float r_min = fminf(fminf(c0, c1), fminf(c2, c3));
const float r_max = fmaxf(fmaxf(c0, c1), fmaxf(c2, c3));
const float p0 = t0x * nx + t0y * ny;
const float p1 = t1x * nx + t1y * ny;
const float p2 = t2x * nx + t2y * ny;
const float t_min = fminf(fminf(p0, p1), p2);
const float t_max = fmaxf(fmaxf(p0, p1), p2);
if (r_max < t_min - eps || t_max < r_min - eps) {
output[idx] = 0.0f;
return;
}
}
output[idx] = 1.0f;
}
torch::Tensor sat_rect_triangle_cuda(
torch::Tensor robot_min,
torch::Tensor robot_max,
torch::Tensor env_triangle,
float eps
) {
const int B = robot_min.size(0);
auto output = torch::empty({B, 1}, robot_min.options());
const int threads = 256;
const int blocks = (B + threads - 1) / threads;
sat_rect_triangle_kernel<<<blocks, threads>>>(
robot_min.data_ptr<float>(),
robot_max.data_ptr<float>(),
env_triangle.data_ptr<float>(),
output.data_ptr<float>(),
eps,
B
);
return output;
}
"""
sat_module = load_inline(
name="sat_rect_triangle_v4",
cpp_sources="",
cuda_sources=cuda_source,
functions=["sat_rect_triangle_cuda"],
extra_cuda_cflags=["-O3", "-use_fast_math"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, eps: float = 1e-6):
super().__init__()
self.eps = eps
def forward(
self,
robot_min: torch.Tensor,
robot_max: torch.Tensor,
env_triangle: torch.Tensor,
) -> torch.Tensor:
robot_min = robot_min.contiguous()
robot_max = robot_max.contiguous()
env_triangle = env_triangle.contiguous()
return sat_module.sat_rect_triangle_cuda(robot_min, robot_max, env_triangle, self.eps)
```
This should work well. The kernel:
1. Uses early returns for better branch prediction
2. Computes everything in registers
3. Uses __restrict__ for memory access optimization
4. Matches the reference logic exactly
Let me submit this implementation.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void sat_rect_triangle_kernel(\n const float* __restrict__ robot_min,\n const float* __restrict__ robot_max,\n const float* __restrict__ env_triangle,\n float* __restrict__ output,\n float eps,\n int B\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= B) return;\n \n // Load rectangle bounds (contiguous row-major)\n const float x0 = robot_min[idx * 2];\n const float y0 = robot_min[idx * 2 + 1];\n const float x1 = robot_max[idx * 2];\n const float y1 = robot_max[idx * 2 + 1];\n \n // Load triangle vertices (B, 3, 2) contiguous\n const float t0x = env_triangle[idx * 6];\n const float t0y = env_triangle[idx * 6 + 1];\n const float t1x = env_triangle[idx * 6 + 2];\n const float t1y = env_triangle[idx * 6 + 3];\n const float t2x = env_triangle[idx * 6 + 4];\n const float t2y = env_triangle[idx * 6 + 5];\n \n // Test X axis [1, 0]\n const float t_x_min = fminf(fminf(t0x, t1x), t2x);\n const float t_x_max = fmaxf(fmaxf(t0x, t1x), t2x);\n if (x1 < t_x_min - eps || t_x_max < x0 - eps) {\n output[idx] = 0.0f;\n return;\n }\n \n // Test Y axis [0, 1]\n const float t_y_min = fminf(fminf(t0y, t1y), t2y);\n const float t_y_max = fmaxf(fmaxf(t0y, t1y), t2y);\n if (y1 < t_y_min - eps || t_y_max < y0 - eps) {\n output[idx] = 0.0f;\n return;\n }\n \n // Edge 0: v0 -> v1, normal = perpendicular (rotated 90 CCW)\n {\n const float ex = t1x - t0x;\n const float ey = t1y - t0y;\n const float len = sqrtf(ex * ex + ey * ey) + 1e-8f;\n const float nx = -ey / len;\n const float ny = ex / len;\n \n const float c0 = x0 * nx + y0 * ny;\n const float c1 = x1 * nx + y0 * ny;\n const float c2 = x1 * nx + y1 * ny;\n const float c3 = x0 * nx + y1 * ny;\n const float r_min = fminf(fminf(c0, c1), fminf(c2, c3));\n const float r_max = fmaxf(fmaxf(c0, c1), fmaxf(c2, c3));\n \n const float p0 = t0x * nx + t0y * ny;\n const float p1 = t1x * nx + t1y * ny;\n const float p2 = t2x * nx + t2y * ny;\n const float t_min = fminf(fminf(p0, p1), p2);\n const float t_max = fmaxf(fmaxf(p0, p1), p2);\n \n if (r_max < t_min - eps || t_max < r_min - eps) {\n output[idx] = 0.0f;\n return;\n }\n }\n \n // Edge 1: v1 -> v2\n {\n const float ex = t2x - t1x;\n const float ey = t2y - t1y;\n const float len = sqrtf(ex * ex + ey * ey) + 1e-8f;\n const float nx = -ey / len;\n const float ny = ex / len;\n \n const float c0 = x0 * nx + y0 * ny;\n const float c1 = x1 * nx + y0 * ny;\n const float c2 = x1 * nx + y1 * ny;\n const float c3 = x0 * nx + y1 * ny;\n const float r_min = fminf(fminf(c0, c1), fminf(c2, c3));\n const float r_max = fmaxf(fmaxf(c0, c1), fmaxf(c2, c3));\n \n const float p0 = t0x * nx + t0y * ny;\n const float p1 = t1x * nx + t1y * ny;\n const float p2 = t2x * nx + t2y * ny;\n const float t_min = fminf(fminf(p0, p1), p2);\n const float t_max = fmaxf(fmaxf(p0, p1), p2);\n \n if (r_max < t_min - eps || t_max < r_min - eps) {\n output[idx] = 0.0f;\n return;\n }\n }\n \n // Edge 2: v2 -> v0\n {\n const float ex = t0x - t2x;\n const float ey = t0y - t2y;\n const float len = sqrtf(ex * ex + ey * ey) + 1e-8f;\n const float nx = -ey / len;\n const float ny = ex / len;\n \n const float c0 = x0 * nx + y0 * ny;\n const float c1 = x1 * nx + y0 * ny;\n const float c2 = x1 * nx + y1 * ny;\n const float c3 = x0 * nx + y1 * ny;\n const float r_min = fminf(fminf(c0, c1), fminf(c2, c3));\n const float r_max = fmaxf(fmaxf(c0, c1), fmaxf(c2, c3));\n \n const float p0 = t0x * nx + t0y * ny;\n const float p1 = t1x * nx + t1y * ny;\n const float p2 = t2x * nx + t2y * ny;\n const float t_min = fminf(fminf(p0, p1), p2);\n const float t_max = fmaxf(fmaxf(p0, p1), p2);\n \n if (r_max < t_min - eps || t_max < r_min - eps) {\n output[idx] = 0.0f;\n return;\n }\n }\n \n output[idx] = 1.0f;\n}\n\ntorch::Tensor sat_rect_triangle_cuda(\n torch::Tensor robot_min,\n torch::Tensor robot_max,\n torch::Tensor env_triangle,\n float eps\n) {\n const int B = robot_min.size(0);\n auto output = torch::empty({B, 1}, robot_min.options());\n \n const int threads = 256;\n const int blocks = (B + threads - 1) / threads;\n \n sat_rect_triangle_kernel<<<blocks, threads>>>(\n robot_min.data_ptr<float>(),\n robot_max.data_ptr<float>(),\n env_triangle.data_ptr<float>(),\n output.data_ptr<float>(),\n eps,\n B\n );\n \n return output;\n}\n\"\"\"\n\nsat_module = load_inline(\n name=\"sat_rect_triangle\",\n cpp_sources=\"\",\n cuda_sources=cuda_source,\n functions=[\"sat_rect_triangle_cuda\"],\n extra_cuda_cflags=[\"-O3\", \"-use_fast_math\"],\n verbose=False,\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, eps: float = 1e-6):\n super().__init__()\n self.eps = eps\n \n def forward(\n self,\n robot_min: torch.Tensor,\n robot_max: torch.Tensor,\n env_triangle: torch.Tensor,\n ) -> torch.Tensor:\n robot_min = robot_min.contiguous()\n robot_max = robot_max.contiguous()\n env_triangle = env_triangle.contiguous()\n return sat_module.sat_rect_triangle_cuda(robot_min, robot_max, env_triangle, self.eps)"
}