Fix NCCL 2.21.5 compatibility issues

PiperOrigin-RevId: 623398701
tensorflow · Apr 10, 2024 · 6b70fa2 · 6b70fa2
1 parent 25f4020
commit 6b70fa2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/third_party/xla/xla/service/gpu/mock_nccl_xml.cc b/third_party/xla/xla/service/gpu/mock_nccl_xml.cc
@@ -50,8 +50,9 @@ class NcclTopoXmlParser {
   Status Parse(ncclXmlNode* head) {
     if (head && head->type == NODE_TYPE_SINGLE) return absl::OkStatus();
     while (true) {
-      if (xml_->maxIndex == MAX_NODES) {
-        return absl::InternalError("XML parser is limited to 1024 nodes");
+      if (xml_->maxIndex == xml_->maxNodes) {
+        return absl::InternalError(absl::StrFormat(
+            "XML parser is limited to %ld nodes", xml_->maxNodes));
       }
       ncclXmlNode* node = xml_->nodes + xml_->maxIndex;
       memset(node, 0, sizeof(ncclXmlNode));

diff --git a/third_party/xla/xla/service/gpu/mock_nccl_xml_test.cc b/third_party/xla/xla/service/gpu/mock_nccl_xml_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/service/gpu/mock_nccl_xml.h"
 
+#include <cstdlib>
 #include <memory>
 #include <string>
 
@@ -41,8 +42,11 @@ TEST_F(MockNcclXmlParserTest, PciNic) {
       </nic>
     </pci>
   )";
-  auto xml = std::make_unique<ncclXml>();
-  auto result = MockTopoGetXml(original, xml.get());
+
+  ncclXml *xml;
+  xmlAlloc(&xml, 1024);
+  std::unique_ptr<void, void (*)(void *)> xml_ptr(xml, free);
+  auto result = MockTopoGetXml(original, xml);
 
   EXPECT_EQ(OkStatus(), result);
   EXPECT_EQ(xml->maxIndex, 3);
@@ -74,8 +78,10 @@ TEST_F(MockNcclXmlParserTest, GpuNvlink) {
       <nvlink target="0000:c7:00.0" count="2" tclass="0x068000"/>
     </gpu>
   )";
-  auto xml = std::make_unique<ncclXml>();
-  auto result = MockTopoGetXml(original, xml.get());
+  ncclXml *xml;
+  xmlAlloc(&xml, 1024);
+  std::unique_ptr<void, void (*)(void *)> xml_ptr(xml, free);
+  auto result = MockTopoGetXml(original, xml);
   EXPECT_EQ(OkStatus(), result);
   EXPECT_EQ(xml->maxIndex, 2);
   EXPECT_EQ(std::string(xml->nodes[0].name), "gpu");