Pikaliov/DCN_custom_op
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: DCNv4 custom op mirror setup

Browse Source 
- Add enhanced README with project structure and quick start guide
- Initialize repository with DCNv4 CUDA extension (PyTorch module)
- Include classification, detection, and segmentation subdirectories
- Reference upstream OpenGVLab DCNv4 implementation

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
Pikaliov
2026-06-11 10:30:44 +03:00
commit 1b3206b6a7
290 changed files with 41632 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

238
DCNv4_op/scripts/test_dcnv4_strip.py Normal file
View File

@@ -0,0 +1,238 @@
"""Strip-DCN feasibility gate (Phase A3, PLAN_code_stripDCN_stripMix).
Validates that the stock compiled extension runs rectangular kernels:
a (1, 9) / (9, 1) strip hits the existing K=9 template instantiation
(switch-K in dcnv4_im2col_cuda.cuh covers {9, 25, 49}; the kernel body is
generic over kernel_h/kernel_w at runtime).
Checks:
1. fwd correctness vs a pure-PyTorch bilinear reference for (1,9), (9,1)
and (3,3) sanity, fp32;
2. bwd correctness (grads wrt value and offset_mask) vs reference autograd;
3. padding channels of offset_mask (beyond G*K*3) are ignored;
4. fp16 fwd: finite, close to fp32 reference;
5. benchmark: strip (1,9) vs square 7 (K=49) vs square 3 (K=9) at the
SOFIA Stage-1 shape [8, 64*64, 192];
6. DCNv4Strip module smoke (zero-init => zero output; random init finite).
PASS gate: max|d| < 1e-4 fp32 (fwd and grads), fp16 finite, strip fwd+bwd
speedup >= 2.5x vs square 7.
Run on the training server (CUDA build of the extension required):
python scripts/test_dcnv4_strip.py
"""
from __future__ import annotations
import time
import torch
from DCNv4 import DCNv4Strip
from DCNv4.functions import DCNv4Function
ATOL_F32 = 1e-4
RTOL_F16 = 3e-2
def ref_dcnv4(
value: torch.Tensor, # [N, H, W, G*D]
offset_mask: torch.Tensor, # [N, H, W, P] (P >= G*K*3, tail = padding)
kh: int, kw: int,
pad_h: int, pad_w: int,
group: int,
offset_scale: float = 1.0,
) -> torch.Tensor:
"""Pure-PyTorch reference mirroring dcnv4_im2col_cuda.cuh semantics.
Point order: outer loop over i in [0, kw), inner over j in [0, kh)
(m = i*kh + j). Per group g the slab offset_mask[..., g*K*3:(g+1)*K*3]
holds K interleaved (dx, dy) pairs first, then K mask values (no softmax).
Sampling location for output pixel (hi, wi):
x = wi - pad_w + (i + dx) * offset_scale_adj, analogous for y,
which for stride=1, dilation=1 reduces to x = wi - pad_w + i + dx
when offset_scale == 1 (matches p0_w_/p0_h_ algebra in the kernel).
Bilinear sampling, zero outside the image. Differentiable.
"""
n, h, w, c = value.shape
d = c // group
k = kh * kw
dev = value.dtype
# Point grid in CUDA order: m = i*kh + j.
ii = torch.arange(kw, device=value.device).repeat_interleave(kh) # [K]
jj = torch.arange(kh, device=value.device).repeat(kw) # [K]
base_y = torch.arange(h, device=value.device).view(1, h, 1, 1)
base_x = torch.arange(w, device=value.device).view(1, 1, w, 1)
out = value.new_zeros(n, h, w, group, d)
half_w = (kw - 1) // 2
half_h = (kh - 1) // 2
for g in range(group):
slab = offset_mask[..., g * k * 3: (g + 1) * k * 3]
offs = slab[..., : k * 2].reshape(n, h, w, k, 2)
mask = slab[..., k * 2: k * 3] # [N,H,W,K]
dx, dy = offs[..., 0], offs[..., 1] # [N,H,W,K]
# p0 - centre*scale + (i*dil + dx)*scale (stride=1, dil=1)
x = (base_x - pad_w + half_w).to(dev) \
+ ((ii.view(1, 1, 1, k) - half_w).to(dev) + dx) * offset_scale
y = (base_y - pad_h + half_h).to(dev) \
+ ((jj.view(1, 1, 1, k) - half_h).to(dev) + dy) * offset_scale
x0 = torch.floor(x); y0 = torch.floor(y)
wx1 = x - x0; wy1 = y - y0
wx0 = 1.0 - wx1; wy0 = 1.0 - wy1
vg = value[..., g * d: (g + 1) * d] # [N,H,W,D]
acc = value.new_zeros(n, h, w, k, d)
for oy, wy_ in ((y0, wy0), (y0 + 1, wy1)):
for ox, wx_ in ((x0, wx0), (x0 + 1, wx1)):
inside = (oy >= 0) & (oy <= h - 1) & (ox >= 0) & (ox <= w - 1)
oy_c = oy.clamp(0, h - 1).long()
ox_c = ox.clamp(0, w - 1).long()
# Gather vg at [n, oy_c, ox_c] for every (h, w, k).
idx = (oy_c * w + ox_c).reshape(n, -1) # [N, H*W*K]
flat = vg.reshape(n, h * w, d)
samp = torch.gather(
flat, 1, idx.unsqueeze(-1).expand(-1, -1, d),
).reshape(n, h, w, k, d)
wgt = (wy_ * wx_ * inside.to(dev)).unsqueeze(-1)
acc = acc + samp * wgt
out[..., g, :] = (acc * mask.unsqueeze(-1)).sum(dim=3)
return out.reshape(n, h * w, c)
def run_op(value_seq, offset_mask, kh, kw, pad_h, pad_w, group, gc):
return DCNv4Function.apply(
value_seq.view(value_seq.shape[0], int(value_seq.shape[1] ** 0.5),
int(value_seq.shape[1] ** 0.5), -1),
offset_mask,
kh, kw, 1, 1, pad_h, pad_w, 1, 1,
group, gc, 1.0, 256, 0,
)
def check_case(kh: int, kw: int, n=2, h=16, w=16, c=64, group=4) -> None:
torch.manual_seed(42)
gc = c // group
k = kh * kw
p = ((group * k * 3 + 7) // 8) * 8
pad_h, pad_w = (kh - 1) // 2, (kw - 1) // 2
value = torch.randn(n, h, w, c, device="cuda", dtype=torch.float32)
om = torch.zeros(n, h, w, p, device="cuda", dtype=torch.float32)
om[..., : group * k * 3] = torch.randn(n, h, w, group * k * 3, device="cuda") * 0.7
v1 = value.clone().requires_grad_(True)
o1 = om.clone().requires_grad_(True)
v2 = value.clone().requires_grad_(True)
o2 = om.clone().requires_grad_(True)
out_op = DCNv4Function.apply(
v1, o1, kh, kw, 1, 1, pad_h, pad_w, 1, 1, group, gc, 1.0, 256, 0,
).reshape(n, h * w, c)
out_ref = ref_dcnv4(v2, o2, kh, kw, pad_h, pad_w, group)
d_fwd = (out_op - out_ref).abs().max().item()
assert d_fwd < ATOL_F32, f"({kh}x{kw}) fwd diverges: {d_fwd:.2e}"
g = torch.randn_like(out_op)
out_op.backward(g)
out_ref.backward(g)
d_gv = (v1.grad - v2.grad).abs().max().item()
d_go = (o1.grad[..., : group * k * 3] - o2.grad[..., : group * k * 3]).abs().max().item()
assert d_gv < ATOL_F32, f"({kh}x{kw}) grad_value diverges: {d_gv:.2e}"
assert d_go < ATOL_F32, f"({kh}x{kw}) grad_offset diverges: {d_go:.2e}"
# Padding channels must be ignored: poison them, output must not change.
om_poison = om.clone()
om_poison[..., group * k * 3:] = 1e6
out_poison = DCNv4Function.apply(
value, om_poison, kh, kw, 1, 1, pad_h, pad_w, 1, 1, group, gc, 1.0, 256, 0,
).reshape(n, h * w, c)
d_pad = (out_poison - out_op.detach()).abs().max().item()
assert d_pad == 0.0, f"({kh}x{kw}) padding channels are NOT ignored: {d_pad:.2e}"
# fp16 forward: finite + close to fp32 reference.
out_h = DCNv4Function.apply(
value.half(), om.half(), kh, kw, 1, 1, pad_h, pad_w, 1, 1,
group, gc, 1.0, 256, 0,
).reshape(n, h * w, c)
assert torch.isfinite(out_h).all(), f"({kh}x{kw}) fp16 produced non-finite"
rel = ((out_h.float() - out_ref).abs() / (out_ref.abs() + 1.0)).max().item()
assert rel < RTOL_F16, f"({kh}x{kw}) fp16 rel err {rel:.2e}"
print(f" OK ({kh}x{kw}): fwd {d_fwd:.1e}, dval {d_gv:.1e}, "
f"doff {d_go:.1e}, pad ignored, fp16 rel {rel:.1e}")
def bench(kh: int, kw: int, label: str, n=8, hw=64, c=192, group=12,
iters=50) -> float:
gc = c // group
k = kh * kw
p = ((group * k * 3 + 7) // 8) * 8
value = torch.randn(n, hw, hw, c, device="cuda", requires_grad=True)
om = torch.randn(n, hw, hw, p, device="cuda", requires_grad=True)
args = (kh, kw, 1, 1, (kh - 1) // 2, (kw - 1) // 2, 1, 1, group, gc, 1.0, 256, 0)
for _ in range(10): # warmup
out = DCNv4Function.apply(value, om, *args)
out.sum().backward()
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(iters):
out = DCNv4Function.apply(value, om, *args)
out.sum().backward()
torch.cuda.synchronize()
ms = (time.perf_counter() - t0) / iters * 1e3
print(f" {label:14s} K={k:2d}: {ms:7.3f} ms/iter (fwd+bwd, "
f"[{n},{hw}x{hw},{c}], G={group})")
return ms
def main() -> None:
assert torch.cuda.is_available(), "CUDA required"
print("== correctness ==")
check_case(3, 3) # sanity: known-good square path (K=9)
check_case(1, 9) # strip-H (K=9, reuses square-3 instantiation)
check_case(9, 1) # strip-V
check_case(1, 9, c=192, group=12) # SOFIA Stage-1 width
print("== DCNv4Strip module ==")
m = DCNv4Strip(channels=192, k=9, orientation="h", group=12).cuda()
assert (m.kernel_h, m.kernel_w) == (1, 9), "orientation='h' must give (1,9)"
mv = DCNv4Strip(channels=192, k=9, orientation="v", group=12).cuda()
assert (mv.kernel_h, mv.kernel_w) == (9, 1), "orientation='v' must give (9,1)"
x = torch.randn(2, 64 * 64, 192, device="cuda")
y = m(x, shape=(64, 64))
assert y.shape == x.shape
assert y.abs().max().item() < 1e-6, "zero-init must give ~zero output"
torch.nn.init.normal_(m.offset_mask.weight, std=0.02)
torch.nn.init.normal_(m.offset_mask.bias, std=0.02)
x = x.requires_grad_(True)
y = m(x, shape=(64, 64))
assert torch.isfinite(y).all() and y.abs().max().item() > 0
y.sum().backward()
assert x.grad is not None and torch.isfinite(x.grad).all()
print(" OK: module smoke (orientations, zero-init ~0, random init finite, bwd)")
print("== benchmark (SOFIA Stage-1 shape) ==")
t_strip = bench(1, 9, "strip (1,9)")
t_sq3 = bench(3, 3, "square 3")
t_sq7 = bench(7, 7, "square 7")
speedup = t_sq7 / t_strip
print(f" strip vs square7 fwd+bwd speedup: {speedup:.2f}x "
f"(target >= 2.5x, hard gate >= 1.8x), vs square3: {t_sq3 / t_strip:.2f}x")
if speedup < 2.5:
print(f" WARNING: speedup {speedup:.2f}x below 2.5x target "
f"(noisy benchmark? rerun on idle GPU)")
assert speedup >= 1.8, "GATE FAILED: strip speedup below hard gate 1.8x"
print("\nALL GATES PASSED — strip-DCN is viable on the stock binary.")
if __name__ == "__main__":
main()