[Mlir-commits] [mlir] [mlir][gpu][nvptx] Remove null terminator when outputting PTX (PR #133019)

[llvmlistbot at llvm.org]

https://github.com/modiking updated https://github.com/llvm/llvm-project/pull/133019

>From 5317a2a4bc08b04d1f2929b154c22aece40d1d4d Mon Sep 17 00:00:00 2001
From: Modi Mo <mmo at nvidia.com>
Date: Mon, 24 Mar 2025 15:45:44 -0700
Subject: [PATCH 1/3] Remove null terminator when outputting PTX

---
 mlir/lib/Target/LLVM/NVVM/Target.cpp               | 7 ++-----
 mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp | 5 ++++-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 023c65b3dd9df..2e7f5353df244 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -723,11 +723,8 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
 
   // Return PTX if the compilation target is `assembly`.
   if (targetOptions.getCompilationTarget() ==
-      gpu::CompilationTarget::Assembly) {
-    // Make sure to include the null terminator.
-    StringRef bin(serializedISA->c_str(), serializedISA->size() + 1);
-    return SmallVector<char, 0>(bin.begin(), bin.end());
-  }
+      gpu::CompilationTarget::Assembly)
+    return SmallVector<char, 0>(serializedISA->begin(), serializedISA->end());
 
   std::optional<SmallVector<char, 0>> result;
   moduleToObjectTimer.startTimer();
diff --git a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
index eabfd1c4d32eb..1231019e5dadf 100644
--- a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
@@ -129,7 +129,10 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {
     ASSERT_TRUE(!object->empty());
 
     ASSERT_TRUE(
-        StringRef(object->data(), object->size()).contains("nvvm_kernel"));
+        StringRef(object->data(), object->size()).contains("nvvm_kernel"));     
+    ASSERT_TRUE(
+        StringRef(object->data(), object->size()).count('\0') == 0);
+
   }
 }
 

>From 616df37ce2ac486aaea1a20e3143f3c41c4fec28 Mon Sep 17 00:00:00 2001
From: Modi Mo <mmo at nvidia.com>
Date: Tue, 25 Mar 2025 16:48:48 -0700
Subject: [PATCH 2/3] clang-format

---
 mlir/lib/Target/LLVM/NVVM/Target.cpp               | 3 +--
 mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp | 6 ++----
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 2e7f5353df244..586748df8d154 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -722,8 +722,7 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
 #undef DEBUG_TYPE
 
   // Return PTX if the compilation target is `assembly`.
-  if (targetOptions.getCompilationTarget() ==
-      gpu::CompilationTarget::Assembly)
+  if (targetOptions.getCompilationTarget() == gpu::CompilationTarget::Assembly)
     return SmallVector<char, 0>(serializedISA->begin(), serializedISA->end());
 
   std::optional<SmallVector<char, 0>> result;
diff --git a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
index 1231019e5dadf..cae713a1ce1d2 100644
--- a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
@@ -129,10 +129,8 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {
     ASSERT_TRUE(!object->empty());
 
     ASSERT_TRUE(
-        StringRef(object->data(), object->size()).contains("nvvm_kernel"));     
-    ASSERT_TRUE(
-        StringRef(object->data(), object->size()).count('\0') == 0);
-
+        StringRef(object->data(), object->size()).contains("nvvm_kernel"));
+    ASSERT_TRUE(StringRef(object->data(), object->size()).count('\0') == 0);
   }
 }
 

>From c2dad87d802137748d211e2bef2deebe8e98d88c Mon Sep 17 00:00:00 2001
From: Modi Mo <mmo at nvidia.com>
Date: Wed, 2 Apr 2025 23:00:25 -0700
Subject: [PATCH 3/3] add null when assembly gets embedded into the binary

---
 mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
index ade239c526af8..8d4a0bcf8adbf 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
@@ -116,8 +116,11 @@ LogicalResult SelectObjectAttrImpl::embedBinary(
   llvm::Module *module = moduleTranslation.getLLVMModule();
 
   // Embed the object as a global string.
+  // Add null for assembly output for JIT paths that expect null-terminated
+  // strings.
+  bool addNull = (object.getFormat() == gpu::CompilationTarget::Assembly);
   llvm::Constant *binary = llvm::ConstantDataArray::getString(
-      builder.getContext(), object.getObject().getValue(), false);
+      builder.getContext(), object.getObject().getValue(), addNull);
   llvm::GlobalVariable *serializedObj =
       new llvm::GlobalVariable(*module, binary->getType(), true,
                                llvm::GlobalValue::LinkageTypes::InternalLinkage,