PyTorch C++ 扩展：有效地累积张量中线程的结果

Question

我正在实现一个采用并行性的 PyTorch C++ 扩展。所需的结果是一个 N x D 矩阵，每个线程计算该结果的一行。累积这些行的推荐方法是什么？

---

为了将其放在上下文中，我的代码如下所示（简化）：

at::Tensor nearestKKeys(at::Tensor queries, at::Tensor keys, int k, int maxLeafSize) {
    /*
    queries: (Nq, D)
    keys: (Nk, D)
    k: int
    return: (Nq, k)
    */
    int Nq = queries.size(0);

    at::Tensor result = at::empty({Nq, k}, at::ScalarType::Int);
    at::Tensor indices = at::arange({Nk}, keys.device())
    BallTreePtr pBallTree = buildBallTree(
        keys,
        indices,
        maxLeafSize,
        0
    );
    
    #pragma omp parallel for
    for (int n=0; n<Nq; n++) {
        at::Tensor query = queries.index({n});
        
        BestMatchesPtr pBestMatches = std::make_shared<BestMatches>(k);
        // places the desired result for this query in *pBestMatches
        pBallTree->query(query, query_norm, pBestMatches, k);
        
        // place the desired row in results
        at::Tensor matches = pBestMatches->getMatches();
        result.index_put_({n}, matches);
   }
}

目前这效率不是很高，我的分析器告诉我线程花费了大量时间相互等待。我相信这是因为

result

变量在线程之间共享。我不认为

pBallTree

被共享真的很重要（因为它是一个指针），但是请告诉我我是否弄错了。因此，我的问题是：累积这些行的推荐方法是什么？

Answer 1

分配通常不能很好地扩展。问题是，您为每次迭代执行

std::make_shared<BestMatches>

，它不仅需要分配一个

BestMatches

对象，还需要分配一个 控制块（引用计数数据结构）。这意味着每次迭代至少进行 2 次分配。

我不确定你是否真的需要一个共享指针。一个简单的

std::unique_ptr

在这里当然就足够了，因为指针实际上并不共享。事实上，我认为你根本不需要动态分配。事实上，

BestMatches

当然可以在并行部分中在堆栈上分配一次，并在多次迭代之间重用。如果您需要调用构造函数并且无法轻松回收该对象，那么您可以使用 placement new 回收该内存区域（不过不要忘记调用析构函数）。

我想知道像

at::Tensor matches = pBestMatches->getMatches();

这样的行是否也执行分配。如果

at::Tensor

不是张量的视图，那么应该分配内存并且肯定会完成副本。这两种操作都无法扩展。不能使用视图或避免创建副本吗？

我假设循环的其他功能是线程安全的。竞争条件不仅会破坏代码，还会使代码变慢。

我还假设

result.index_put_

只是执行基本的一维张量复制。如果是这样，那么访问

results

并不是真正的问题，因为默认的 OpenMP 调度策略往往不会导致太多错误共享问题。