莖節點與葉子節點

莖節點與葉子節點皆適用KdAccelNode來表示
註意：這裏使用了匿名union

union有個特性:內部類型共用一段內存，且大小為內部最大類型的大小。

struct KdAccelNode {
<KdAccelNode Methods>
    union {
        Float split; // Interior
        int onePrimitive; // Leaf
        int primitiveIndicesOffset; // Leaf
    };
    union {
        int flags; // Both
        int nPrims; // Leaf
        int aboveChild; // Interior
    };
};
 

使用8字節存儲葉子節點（2個union各4byte），其中flags中的低二位2bit用於區分分割方向以及葉子節點，nPrims的高30位存儲圖元存儲數。InitInterior函數也是差不多的操作。
primNums數組存儲了圖元容器的index，

void KdAccelNode::InitLeaf(int *primNums, int np,
                           std::vector<int> *primitiveIndices) {
    //標記為葉子
    flags = 3;
    //np為圖元個數，這樣做可以保證上一步flags=3時的存儲的數據不會覆蓋掉
    nPrims |= (np << 2);
    //存儲圖元index，如果是0~1個，使用onePrimitive來存儲
    if (np == 0)
        onePrimitive = 0;
    else if (np == 1)
        onePrimitive = primNums[0];
    else {
    //多個圖元的情況下，使用存儲index數組的偏移值，之後存儲每個圖元index
        primitiveIndicesOffset = primitiveIndices->size();
        for (int i = 0; i < np; ++i) primitiveIndices->push_back(primNums[i]);
    }
}
 

構建樹

KdTreeAccel::KdTreeAccel(std::vector<std::shared_ptr<Primitive>> p,
                         int isectCost, int traversalCost, Float emptyBonus,
                         int maxPrims, int maxDepth)
    : isectCost(isectCost),
      traversalCost(traversalCost),
      maxPrims(maxPrims),
      emptyBonus(emptyBonus),
      primitives(std::move(p)) {
    //nextFreeNode為下一個可用節點的index，nAllocedNodes為分配空間的節點總數
    //將兩者設為0，可以讓系統在初始化樹中的第一個節點時分配內存而不是在啟動就分配內存。
    //這兩個變量用於判斷在此次遞歸中是否需要申請內存塊
    ProfilePhase _(Prof::AccelConstruction);
    nextFreeNode = nAllocedNodes = 0;
    //計算深度的經驗公式:8 + 1.3 log(N)
    if (maxDepth <= 0)
        maxDepth = std::round(8 + 1.3f * Log2Int(int64_t(primitives.size())));

    //計算kd樹結構的邊界盒，以及創建各圖元邊界盒數組
    std::vector<Bounds3f> primBounds;
    primBounds.reserve(primitives.size());
    for (const std::shared_ptr<Primitive> &prim : primitives) {
        Bounds3f b = prim->WorldBound();
        bounds = Union(bounds, b);
        primBounds.push_back(b);
    }

    //分配內存的相應變量
    
    //edges用於存儲x，y，z三個方向上的所有可能切割圖元位置
    //因為邊界盒有min、max值，所有這裏的需要圖元個數*2
    std::unique_ptr<BoundEdge[]> edges[3];
    for (int i = 0; i < 3; ++i)
        edges[i].reset(new BoundEdge[2 * primitives.size()]);
    //以最壞打算確定prim1的空間
    std::unique_ptr<int[]> prims0(new int[primitives.size()]);
    std::unique_ptr<int[]> prims1(new int[(maxDepth + 1) * primitives.size()]);

    //primNums 為與當前節點重疊的圖元index數組
    std::unique_ptr<int[]> primNums(new int[primitives.size()]);
    for (size_t i = 0; i < primitives.size(); ++i) primNums[i] = i;

    //開始遞歸構建kd樹
    buildTree(0, bounds, primBounds, primNums.get(), primitives.size(),
              maxDepth, edges, prims0.get(), prims1.get());
}
 

nodeNum為創建節點的偏移量（primNums中的index）

void KdTreeAccel::buildTree(int nodeNum, const Bounds3f &nodeBounds,
                            const std::vector<Bounds3f> &allPrimBounds,
                            int *primNums, int nPrimitives, int depth,
                            const std::unique_ptr<BoundEdge[]> edges[3],
                            int *prims0, int *prims1, int badRefines) {
    CHECK_EQ(nodeNum, nextFreeNode);
    //當分配的內存使用完時，分配下一塊內存
    if (nextFreeNode == nAllocedNodes) {
        int nNewAllocNodes = std::max(2 * nAllocedNodes, 512);
        KdAccelNode *n = AllocAligned<KdAccelNode>(nNewAllocNodes);
        if (nAllocedNodes > 0) {
            memcpy(n, nodes, nAllocedNodes * sizeof(KdAccelNode));
            FreeAligned(nodes);
        }
        nodes = n;
        nAllocedNodes = nNewAllocNodes;
    }
    ++nextFreeNode;

    //當需要分割的圖元數足夠小或是樹的深度達到一定次數則生成葉子節點停止遞歸
    if (nPrimitives <= maxPrims || depth == 0) {
        nodes[nodeNum].InitLeaf(primNums, nPrimitives, &primitiveIndices);
        return;
    }

    //莖節點需要通過分割來生成，這裏采用了類似之前生成BVH的SAH算法

    int bestAxis = -1, bestOffset = -1;
    Float bestCost = Infinity;
    Float oldCost = isectCost * Float(nPrimitives);
    Float totalSA = nodeBounds.SurfaceArea();
    Float invTotalSA = 1 / totalSA;
    Vector3f d = nodeBounds.pMax - nodeBounds.pMin;

    //選擇一個用於分割圖元的軸向
    //優先選擇具有最大空間性的軸向
    int axis = nodeBounds.MaximumExtent();
    int retries = 0;
    //goto語句跳轉用
retrySplit:

    //初始化對應edges二維數組中對應軸的數組
    //通過圖元數找到對應的index，在通過index獲得邊界盒
    for (int i = 0; i < nPrimitives; ++i) {
        int pn = primNums[i];
        const Bounds3f &bounds = allPrimBounds[pn];
        edges[axis][2 * i] = BoundEdge(bounds.pMin[axis], pn, true);
        edges[axis][2 * i + 1] = BoundEdge(bounds.pMax[axis], pn, false);
    }

    //以切割位置為準進行升序排序，如果位置相同，以類型start為先
    std::sort(&edges[axis][0], &edges[axis][2 * nPrimitives],
              [](const BoundEdge &e0, const BoundEdge &e1) -> bool {
                  if (e0.t == e1.t)
                      return (int)e0.type < (int)e1.type;
                  else
                      return e0.t < e1.t;
              });

    //遍歷所有分割位置來計算最佳分割點
    int nBelow = 0, nAbove = nPrimitives;
    for (int i = 0; i < 2 * nPrimitives; ++i) {
        if (edges[axis][i].type == EdgeType::End) --nAbove;
        //取得分割位置
        Float edgeT = edges[axis][i].t;
        //判斷當前分割點是否在節點的邊界盒內
        if (edgeT > nodeBounds.pMin[axis] && edgeT < nodeBounds.pMax[axis]) {
            // 取得另兩個軸，用於計算分割後2個邊界盒的面積，之後再用公式計算消耗
            int otherAxis0 = (axis + 1) % 3, otherAxis1 = (axis + 2) % 3;
            Float belowSA = 2 * (d[otherAxis0] * d[otherAxis1] +
                                 (edgeT - nodeBounds.pMin[axis]) *
                                     (d[otherAxis0] + d[otherAxis1]));
            Float aboveSA = 2 * (d[otherAxis0] * d[otherAxis1] +
                                 (nodeBounds.pMax[axis] - edgeT) *
                                     (d[otherAxis0] + d[otherAxis1]));
            Float pBelow = belowSA * invTotalSA;
            Float pAbove = aboveSA * invTotalSA;
            Float eb = (nAbove == 0 || nBelow == 0) ? emptyBonus : 0;
            Float cost =
                traversalCost +
                isectCost * (1 - eb) * (pBelow * nBelow + pAbove * nAbove);

            //如果消耗比之前的變量小，則更新對應的變量
            if (cost < bestCost) {
                bestCost = cost;
                bestAxis = axis;
                bestOffset = i;
            }
        }
        if (edges[axis][i].type == EdgeType::Start) ++nBelow;
    }
    CHECK(nBelow == nPrimitives && nAbove == 0);

    //如果沒有找到最優切割位置，則切換軸向繼續尋找
    if (bestAxis == -1 && retries < 2) {
        ++retries;
        axis = (axis + 1) % 3;
        goto retrySplit;
    }
    if (bestCost > oldCost) ++badRefines;
    //消耗大於一定數量 && 圖元數小於一定量 或者無法找到最優切割位置，則直接生成葉子節點，結束遞歸
    //無法找到最優切割位置參考書295頁的圖，即所有圖元都重疊在一起
    //當然可能出現好幾次循環也找不到最佳結果的情況，所以系統給予4次機會，讓其繼續尋找
    if ((bestCost > 4 * oldCost && nPrimitives < 16) || bestAxis == -1 ||
        badRefines == 3) {
        nodes[nodeNum].InitLeaf(primNums, nPrimitives, &primitiveIndices);
        return;
    }

    //將分割的圖元分組，用於下一次遞歸
    int n0 = 0, n1 = 0;
    for (int i = 0; i < bestOffset; ++i)
        if (edges[bestAxis][i].type == EdgeType::Start)
            prims0[n0++] = edges[bestAxis][i].primNum;
    for (int i = bestOffset + 1; i < 2 * nPrimitives; ++i)
        if (edges[bestAxis][i].type == EdgeType::End)
            prims1[n1++] = edges[bestAxis][i].primNum;

    //進行下一次遞歸
    Float tSplit = edges[bestAxis][bestOffset].t;
    Bounds3f bounds0 = nodeBounds, bounds1 = nodeBounds;
    //以切割位置為準，重新計算兩個孩子節點的邊界盒
    bounds0.pMax[bestAxis] = bounds1.pMin[bestAxis] = tSplit;
    buildTree(nodeNum + 1, bounds0, allPrimBounds, prims0, n0, depth - 1, edges,
              prims0, prims1 + nPrimitives, badRefines);
    int aboveChild = nextFreeNode;
    nodes[nodeNum].InitInterior(bestAxis, aboveChild, tSplit);
    buildTree(aboveChild, bounds1, allPrimBounds, prims1, n1, depth - 1, edges,
              prims0, prims1 + nPrimitives, badRefines);
}

bool KdTreeAccel::Intersect(const Ray &ray, SurfaceInteraction *isect) const {
    ProfilePhase p(Prof::AccelIntersect);
    Float tMin, tMax;
    if (!bounds.IntersectP(ray, &tMin, &tMax)) {
        return false;
    }

    Vector3f invDir(1 / ray.d.x, 1 / ray.d.y, 1 / ray.d.z);
    PBRT_CONSTEXPR int maxTodo = 64;
    //KdToDo記錄尚未被處理的節點
    KdToDo todo[maxTodo];
    int todoPos = 0;
    bool hit = false;
    const KdAccelNode *node = &nodes[0];
    while (node != nullptr) {
        if (ray.tMax < tMin) break;
        //如果不是葉子節點
        if (!node->IsLeaf()) {

            //首先與莖節點的分割平面做相交測試，以此來確定光線與節點的相交順序
            int axis = node->SplitAxis();
            //射線的對應分量上的t值
            Float tPlane = (node->SplitPos() - ray.o[axis]) * invDir[axis];

            //判斷先後順序取得兩個節點
            const KdAccelNode *firstChild, *secondChild;
            int belowFirst =
                (ray.o[axis] < node->SplitPos()) ||
                (ray.o[axis] == node->SplitPos() && ray.d[axis] <= 0);
            if (belowFirst) {
                firstChild = node + 1;
                secondChild = &nodes[node->AboveChild()];
            } else {
                firstChild = &nodes[node->AboveChild()];
                secondChild = node + 1;
            }

            //t大於max或者為負數代表了射線不會第二個節點相交
            if (tPlane > tMax || tPlane <= 0)
                node = firstChild;
            //小於tmin不會與第一個節點相交
            else if (tPlane < tMin)
                node = secondChild;
            //兩個節點都相交
            else {
                //將第二個節點放入處理隊列
                todo[todoPos].node = secondChild;
                todo[todoPos].tMin = tPlane;
                todo[todoPos].tMax = tMax;
                ++todoPos;
                node = firstChild;
                tMax = tPlane;
            }
        } else {
            //葉子節點則對內部圖元進行相交測試
            int nPrimitives = node->nPrimitives();
            if (nPrimitives == 1) {
                const std::shared_ptr<Primitive> &p =
                    primitives[node->onePrimitive];
                // Check one primitive inside leaf node
                if (p->Intersect(ray, isect)) hit = true;
            } else {
                for (int i = 0; i < nPrimitives; ++i) {
                    int index =
                        primitiveIndices[node->primitiveIndicesOffset + i];
                    const std::shared_ptr<Primitive> &p = primitives[index];
                    // Check one primitive inside leaf node
                    if (p->Intersect(ray, isect)) hit = true;
                }
            }

            //更新下一次遍歷區域的min與max值
            if (todoPos > 0) {
                --todoPos;
                node = todo[todoPos].node;
                tMin = todo[todoPos].tMin;
                tMax = todo[todoPos].tMax;
            } else
                break;
        }
    }
    return hit;
}

結尾

綜上所述，KD樹的分割（每個軸遍歷所有分割方案）與BVH（SAH）分割（每個軸12次分割）更加精確，這樣是的KD樹的求交效率要比BVH好，也導致了渲染時需要花費比BVH更多的時間在構建KD樹上。同時KD樹在求交時會計算與光線與分割面的相交位置是否在[min,max]中來判斷是否需要與莖節點中的兩個節點一一求交，以此來減少求交計算量。而BVH通過深度優先遍歷樹進行求交的。

PBRT筆記(3)——KD樹