I was able to refactor the code to get OverlappingBBoxes mapped over the 100000 testbboxes to go from 44.85 seconds down to 1.3 seconds (a 34.5x speedup). I suspect there still room for improvement.

Code attached this time rather than being pasted in the post.

Improvements:

• I changed the flow of my code a bit:
  • Place bboxes that overlap with the splitting plane into the left node (as opposed to keeping separate).
  • For bboxes in a node, sort them by their splittingdim+1 dimension. So now in OverlappingBBoxes we can binary search and quickly eliminate ~50% bboxes to manually check.

In lieu of proper documentation I found what I believe to be the source code at

FileNameJoin[{$InstallationDirectory, "SystemFiles", "Components", "Compile"}]

Based on what I saw:

• I think Native`UncheckedBlock is similar to "RuntimeOptions" -> "Speed" in Compile.
  • With help from a 2019 WTC presentation I was able to inspect the IR with Compile`CompileToIR and noticed tons of check were omitted when this was used.
  • Not all checks are omitted with this though. NullNode exceptions in binary trees are still checked for example.
  • Is there any other way to turn off more checks?
• Based off the MSE answer here we have some control over what gets inlined.
  • Based on FunctionInlineInformation.m in the source code, we can use the option "InlinePolicy" -> "MaximumCallCount" -> 1 to inline a function that is only called in one spot.
  • I don't know if functions are inlined at a later step, so perhaps this is superfluous.
• list[[i]] seems to avoid bounds checks with Native`UncheckedBlock, but multi arg Part does not.

  • Looks like with UncheckedBlock, unary Part gets converted to

    Native`Unchecked[Compile`GetArrayElement]
    

but multi arg Part does not. * I refactored my AABB node to store a flat list rather than a matrix.

Attachments:

AABBTree.nb

That seems like a valid use of TypeDeclaration. However, I've myself not worked yet with the "Product" type in any significant examples so I can't comment if it is a suitable use for that type.

I'm not sure why it's that bit slower either, but the returned value you get I think is valid. What's failed I think is the attempt to get its full form. This is the same behavior you commented separately with the other questions (point no. 2). I tried to answer it there too

At the moment I only have an answer for two of these questions

• Point no. 2: You can definitely return data structures of any compiled type, from compiled code (as the evaluator doesn't know about compiled types, only about expressions). The front end renders them differently on screen, and nothing seems to work with them outside compiled code. I think the expression you mention with a $Failed is just a failure when attempting to get the full form of the object, but it's perfectly valid for its purpose. I wonder if the failure to show the full form is meaningful and on purpose. Even if it's compiled stuff, you can for example successfully show the full form of a compiled function. I'll try to find out more

• Point no. 6: I think the way to achieve that is with the type "InertExpression". It allows you to put anything inside. On the counterpart, some functions won't work if you apply them to elements of the type "InertExpression", so I understand this is one of the reasons why you wouldn't just make "ListVector"::["InertExpression"] all the time if you know you have uniform types. (There may be more reasons that I'd also like to know)

Maybe someone else could give an answer for the rest of points or I can try to find out more and ask

With the great nudge from @José Antonio Fernández I was able to compile the construction of the AABBTree and refined the OverlappingBBoxes function a bit more. Code posted below.

Example

Initialization:

randomBBox[n_] :=
  With[{llcorners = RandomReal[{-10, 10}, {n, 3}]},
    Transpose[{llcorners, llcorners + RandomReal[{1, 1}, {n, 3}]}, {3, 1, 2}]
  ]

SeedRandom[1];
bboxes = randomBBox[100000];
testbboxes = randomBBox[100000];

Construction used to take 0.235465 seconds (so a 2x speedup):

tree = AABBTree[bboxes]; // AbsoluteTiming
(* {0.107196, Null} *)

Testing all 100000 bboxes used to take 44.8445 seconds (so a 6.5x speedup):

Scan[OverlappingBBoxes[tree, #] &, testbboxes]; // AbsoluteTiming
(* {6.9374, Null} *)

Code

decl = FunctionDeclaration[
    centerNode3D,
    Typed[{"MachineInteger", "PackedArray"::["Real64", 2], "Real64"} -> "BinaryTree"] @ Function[{d, intstagged, x},
       Module[{preamble},
         preamble = ConstantArray[0., 7];
         preamble[[1]] = x;
         preamble[[2]] = Min[intstagged[[All, 2d-1]]];
         preamble[[3]] = Max[intstagged[[All, 2d]]];
         preamble[[4]] = Cast[d, "Real64"];
         CreateDataStructure["BinaryTree", Cast[Append[intstagged, preamble], "InertExpression"]]
       ]
    ]
];

pureAABBTree = Function[{Typed[bboxes, "PackedArray"::["Real64", 3]]},
    Module[{iAABBTree, bboxestagged},
       iAABBTree = Function[{Typed[d, "MachineInteger"], Typed[bboxestagged, "PackedArray"::["Real64", 2]]},
         Module[{len, n, x, mnp, mxp, loc, r=0, l=0, c=0, Sleft, Scent, Sright, node},
          len = Length[bboxestagged];

          n = Ordering[bboxestagged[[All, 2d-1]]][[Quotient[len+1, 2]]];
          x = bboxestagged[[n, 2d-1]];

          If[len <= 100, Return[centerNode3D[d, bboxestagged, x]]];

          mnp = 2d-1;
          mxp = mnp+1;

          loc = ConstantArray[0, Length[bboxestagged]];

          Do[
              If[bboxestagged[[i, mxp]] < x,
                 l++;
                 loc[[i]] = 1,
                 If[bboxestagged[[i, mnp]] <= x,
                   c++;
                   loc[[i]] = 2,
                   r++
                 ]
              ],
              {i, Length[bboxestagged]}
          ];

          Sleft = ConstantArray[0., {l, 7}];
          Scent = ConstantArray[0., {c, 7}];
          Sright = ConstantArray[0., {r, 7}];

          r = l = c = 1;
          Do[
              If[loc[[i]] === 0,
                 Sright[[r++]] = bboxestagged[[i]];,
                 If[loc[[i]] === 1,
                   Sleft[[l++]] = bboxestagged[[i]],
                   Scent[[c++]] = bboxestagged[[i]];
                 ]
              ],
              {i, Length[bboxestagged]}
          ];

          node = centerNode3D[d, Scent, x];

          If[l > 1, BinaryTree`SetLeft[node, iAABBTree[Mod[d+1, 3, 1], Sleft]]];
          If[r > 1, BinaryTree`SetRight[node, iAABBTree[Mod[d+1, 3, 1], Sright]]];

          node
         ]
       ];

       bboxestagged = Table[Append[Flatten[bboxes[[i]]], Cast[i, "Real64"]], {i, Length[bboxes]}][[Ordering[bboxes[[All, 1, 1]]]]];

       iAABBTree[1, bboxestagged]
    ]
];

declbbox = FunctionDeclaration[
    bboxOverlaps, 
    Typed[ConstantArray["Real64", 12] -> "Boolean"] @ Function[
       {epx1, epx2, epy1, epy2, epz1, epz2, x1, x2, y1, y2, z1, z2}, 
       And[
         epx1 <= x1 <= epx2 || epx1 <= x2 <= epx2 || (x1 <= epx1 <= epx2 <= x2),
         epy1 <= y1 <= epy2 || epy1 <= y2 <= epy2 || (y1 <= epy1 <= epy2 <= y2),
         epz1 <= z1 <= epz2 || epz1 <= z2 <= epz2 || (z1 <= epz1 <= epz2 <= z2)
       ]
    ]
];

pureOverlappingBBoxes = Function[
    {
       Typed[tree, "BinaryTree"], 
       Typed[bbox, "PackedArray"::["Real64", 2]]
    },
    Block[{inds, stack, x1, x2, y1, y2, z1, z2, node, a, b, data, x, min, max, d},
       inds = TypeHint[CreateDataStructure["ExtensibleVector"], "ExtensibleVector"::["MachineInteger"]];
       stack = TypeHint[CreateDataStructure["Stack"], "Stack"::["BinaryTree"]];
       stack["Push", tree];

       x1 = bbox[[1, 1]]; x2 = bbox[[1, 2]];
       y1 = bbox[[2, 1]]; y2 = bbox[[2, 2]];
       z1 = bbox[[3, 1]]; z2 = bbox[[3, 2]];

       While[!stack["EmptyQ"],
         node = stack["Pop"];
         data = Cast[node["Data"], "PackedArray"::["Real64", 2]];
         If[Length[data] <= 1, Continue[]];
         x = data[[-1, 1]];
         min = data[[-1, 2]];
         max = data[[-1, 3]];
         d = Round[data[[-1, 4]]];
         a = bbox[[d, 1]];
         b = bbox[[d, 2]];
         If[min <= a <= max || min <= b <= max || (a <= min && max <= b),
          Do[
              If[bboxOverlaps[data[[i, 1]], data[[i, 2]], data[[i, 3]], data[[i, 4]], data[[i, 5]], data[[i, 6]], x1, x2, y1, y2, z1, z2],
                 inds["Append", Round[data[[i, 7]]]];
              ],
              {i, Length[data]-1}
          ]
         ];

         If[a < x && !node["LeftNullQ"], stack["Push", node["Left"]]];
         If[x < b && !node["RightNullQ"], stack["Push", node["Right"]]];
       ];

       inds["Elements"]
    ]
];

AABBTree = 
  FunctionCompile[decl, pureAABBTree, 
   CompilerOptions -> {"AbortHandling" -> False, 
     "LLVMOptimization" -> "ClangOptimization"[3], 
     "OptimizationLevel" -> 3}];

OverlappingBBoxes = 
  FunctionCompile[declbbox, pureOverlappingBBoxes, 
   CompilerOptions -> {"AbortHandling" -> False, 
     "LLVMOptimization" -> "ClangOptimization"[3], 
     "OptimizationLevel" -> 3}];

Hey thanks, I was able to get it to work! I haven't touched anything with parallelization yet, but will look into that soon.

Here's the compiled code:

decl = FunctionDeclaration[
    bboxOverlaps, 
    Typed[ConstantArray["Real64", 12] -> "Boolean"] @ Function[
       {epx1, epx2, epy1, epy2, epz1, epz2, x1, x2, y1, y2, z1, z2}, 
       And[
         epx1 <= x1 <= epx2 || epx1 <= x2 <= epx2 || (x1 <= epx1 <= epx2 <= x2),
         epy1 <= y1 <= epy2 || epy1 <= y2 <= epy2 || (y1 <= epy1 <= epy2 <= y2),
         epz1 <= z1 <= epz2 || epz1 <= z2 <= epz2 || (z1 <= epz1 <= epz2 <= z2)
       ]
    ]
];

pureOverlappingBBoxes = Function[
    {
       Typed[tree, "BinaryTree"], 
       Typed[bbox, "PackedArray"::["Real64", 2]]
    },
    Block[{inds, stack, x1, x2, y1, y2, z1, z2, node, a, b, data, d1, d2, d3, d4, d5, d6, d7, d8, d9, d10, d11},
       inds = TypeHint[CreateDataStructure["ExtensibleVector"], "ExtensibleVector"::["MachineInteger"]];
       stack = TypeHint[CreateDataStructure["Stack"], "Stack"::["BinaryTree"]];
       stack["Push", tree];

       x1 = bbox[[1, 1]]; x2 = bbox[[1, 2]];
       y1 = bbox[[2, 1]]; y2 = bbox[[2, 2]];
       z1 = bbox[[3, 1]]; z2 = bbox[[3, 2]];

       While[!stack["EmptyQ"],
         node = stack["Pop"];
         data = node["Data"];
         d1 = Cast[data[[1]], "Real64"];
         d2 = Cast[data[[2]], "PackedArray"::["Real64", 1]];
         d3 = Cast[data[[3]], "PackedArray"::["Real64", 1]];
         d4 = Cast[data[[4]], "PackedArray"::["Real64", 1]];
         d5 = Cast[data[[5]], "PackedArray"::["Real64", 1]];
         d6 = Cast[data[[6]], "PackedArray"::["Real64", 1]];
         d7 = Cast[data[[7]], "PackedArray"::["Real64", 1]];
         d8 = Cast[data[[8]], "PackedArray"::["MachineInteger", 1]];
         d9 = Cast[data[[9]], "Real64"];
         d10 = Cast[data[[10]], "Real64"];
         d11 = Cast[data[[11]], "MachineInteger"];
         a = bbox[[d11, 1]];
         b = bbox[[d11, 2]];
         If[Length[d2] > 0 && (d9 <= a <= d10 || d9 <= b <= d10 || a <= d9 <= d10 <= b),
          Do[
              If[bboxOverlaps[d2[[i]], d3[[i]], d4[[i]], d5[[i]], d6[[i]], d7[[i]], x1, x2, y1, y2, z1, z2],
                 inds["Append", d8[[i]]];
              ],
              {i, Length[d8]}
          ]
         ];

         If[a < d1 && !node["LeftNullQ"], stack["Push", node["Left"]]];
         If[d1 < b && !node["RightNullQ"], stack["Push", node["Right"]]];
       ];

       inds["Elements"]
    ]
];

(* takes 44 seconds to compile... *)
OverlappingBBoxes = FunctionCompile[decl, pureOverlappingBBoxes];

I also had to modify iAABBTree to make the data in the binary tree a bit more consistent:

iAABBTree[d_, bboxestagged_] :=
    Block[{len, n, x, Sleft, Scent, Sright, node},
       len = Length[bboxestagged];

       n = Ordering[bboxestagged[[All, 2d-1]], {Quotient[len+1, 2]}][[1]];
       x = Mean[bboxestagged[[n, 2d-1 ;; 2d]]];

       If[len <= 16, Return[centerNode3D[d, x, bboxestagged]]];

       {Sleft, Scent, Sright} = partitionIntervals3D[d, bboxestagged, x];

       node = centerNode3D[d, x, Scent];
       If[Length[Sleft] > 0, node["SetLeft", iAABBTree[Mod[d+1, 3, 1], Sleft]];];
       If[Length[Sright] > 0, node["SetRight", iAABBTree[Mod[d+1, 3, 1], Sright]];];

       node
    ]

And now the call that took 44 seconds now takes 11:

SeedRandom[1];
bboxes = randomBBox[100000];
testbboxes = randomBBox[100000];

aabbtree = AABBTree[bboxes]; // AbsoluteTiming
(* {0.262679, Null} *)

Scan[OverlappingBBoxes[aabbtree, #] &, testbboxes]; // AbsoluteTiming
(* {11.5295, Null} *)

Question

Notice how I took all the elements of the binary tree's data and cast them to their appropriate types. I wonder if casting from InertExpression to each type is adding a slowdown. Is there a way to instead convey this information in Typed[tree, "BinaryTree"]?

e.g. something like Typed[tree, "BinaryTree"::[{"Real64", "PackedArray"::["Real64", 1], ...}]]