BallQuery

产品支持情况

产品	是否支持
Ascend 950PR&950DT系列产品	√
Atlas A3系列产品	√
Atlas A2系列产品	√
Atlas 200I/500 A2推理产品	×
Atlas推理系列产品	√
Atlas训练系列产品	√

功能说明

• 算子功能：球查询（Ball Query），PointNet++中的经典算子。对每个查询中心点center_xyz[m,b]，在xyz[b]中按k=0..N-1顺序查找满足距离判定条件的点，收集前sample_num个点的索引。

• 计算公式：

布局：xyz=(B,3,N)，center_xyz=(M,B,3)，idx=(M,B,sample_num)。对每个(m,b)，取中心点$(c_x,c_y,c_z)=center\_xyz[m,b,:]$，计算其与xyz[b]中每个点$(x_k,y_k,z_k)=(xyz[b,0,k],xyz[b,1,k],xyz[b,2,k])$的距离平方：

$$d2=(c_x-x_k{)}^2+(c_y-y_k{)}^2+(c_z-z_k{)}^2$$

距离判定条件：$d2==0$ 或 $min\_radiu{s}^2\le d2<max\_radiu{s}^2$。命中点按k顺序写入idx，首个命中点索引记为first_num，命中数不足sample_num时剩余位置用first_num填充；若某查询球命中数为0，则该球对应的全部sample_num个输出位置均填充0，该结果与命中第0个点在输出上无法区分。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
xyz	输入	所有点的xyz坐标，shape为(B,3,N)，中间维必须为3（坐标在中间维）。	FLOAT、FLOAT16	ND
center_xyz	输入	查询中心点坐标，shape为(M,B,3)，最后一维必须为3，第1维B必须与xyz的第0维一致。	数据类型与xyz保持一致	ND
idx	输出	采样点在xyz中的索引，shape为(M,B,sample_num)。	INT32	ND
min_radius	属性	最小查询半径（环形内边界），必须大于等于0。	Float	-
max_radius	属性	最大查询半径（球外边界），必须大于min_radius。	Float	-
sample_num	属性	每球最大采样点数，必须大于等于1。	Int	-

约束说明

• xyz为3维Tensor，shape为(B,3,N)，中间维必须为3。
• center_xyz为3维Tensor，shape为(M,B,3)，最后一维必须为3。
• xyz的第0维B与center_xyz的第1维B必须一致。
• xyz的第0维B、第2维N，以及center_xyz的第0维M(npoint)取值均不得超过INT32_MAX(2147483647)。
• xyz与center_xyz的dtype必须相同，且为FLOAT或FLOAT16。
• 输出idx的dtype固定为INT32，shape为(M,B,sample_num)。
• 支持零长度维（空张量）输入：当xyz的第2维N=0时，每个查询球均无候选点、命中数为0，idx输出全部为0；当center_xyz的第0维M=0时，输出shape为(0,B,sample_num)；当B=0时，输出shape为(M,0,sample_num)。以上空输入场景均正常执行并返回。
• 支持动态shape：构建图时输入维度可标记为-1（表示该维取值运行时确定），支持未知秩输入（shape标记为-2）；坐标维为3、B一致性、维度不超过INT32_MAX等校验延后至运行时tiling阶段以真实shape执行；秩未知时输出shape推导为(-1,-1,sample_num)。
• min_radius必须为有限值（非NaN、非Inf）且大于等于0。
• max_radius必须为有限值（非NaN、非Inf）且大于0，并必须大于min_radius。
• sample_num必须大于等于1且不超过INT32_MAX(2147483647)。
• 输入Tensor需为连续内存布局（contiguous）：即数据在内存中按行主序紧密排列，不含stride跳步或间隙（概念详见非连续的Tensor）。若传入非连续Tensor（如转置、切片得到的视图），框架会自动连续化处理，产生额外拷贝开销。
• 算子仅支持3D输入，不支持标量、1D或8D场景。

调用说明

调用方式	调用样例	说明
GE图模式	test_geir_ball_query	通过算子IR构图方式调用BallQuery算子，参见算子调用完成算子编译和验证。