| 28 | |
| 29 | |
| 30 | class KNNQuery(Function): |
| 31 | @staticmethod |
| 32 | def forward(ctx, nsample, xyz, new_xyz, offset, new_offset): |
| 33 | """ |
| 34 | input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b) |
| 35 | output: idx: (m, nsample), dist2: (m, nsample) |
| 36 | """ |
| 37 | if new_xyz is None: new_xyz = xyz |
| 38 | assert xyz.is_contiguous() and new_xyz.is_contiguous() |
| 39 | m = new_xyz.shape[0] |
| 40 | idx = torch.cuda.IntTensor(m, nsample).zero_() |
| 41 | dist2 = torch.cuda.FloatTensor(m, nsample).zero_() |
| 42 | pointops_cuda.knnquery_cuda(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2) |
| 43 | return idx, torch.sqrt(dist2) |
| 44 | |
| 45 | knnquery = KNNQuery.apply |
| 46 |
nothing calls this directly
no outgoing calls
no test coverage detected