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

[llvmlistbot at llvm.org]

Author: modiking
Date: 2025-04-03T15:50:54-07:00
New Revision: 9f2feeb1891a4d58823e6cbbb404847d6e43e65b

URL: https://github.com/llvm/llvm-project/commit/9f2feeb1891a4d58823e6cbbb404847d6e43e65b
DIFF: https://github.com/llvm/llvm-project/commit/9f2feeb1891a4d58823e6cbbb404847d6e43e65b.diff

LOG: [mlir][gpu][nvptx] Remove null terminator when outputting PTX (#133019)

PTX source files are expected to only contain ASCII text
(https://docs.nvidia.com/cuda/parallel-thread-execution/#source-format) and no null terminators.

`ptxas` has so far not enforced this but is moving towards doing so.
This revealed a problem where the null terminator is getting printed out
in the output file in MLIR path when outputting ptx directly. Only add the null on the assembly output path for JIT instead of in output of `moduleToObject `.

Added: 
    

Modified: 
    mlir/lib/Target/LLVM/NVVM/Target.cpp
    mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
    mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 023c65b3dd9df..586748df8d154 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -722,12 +722,8 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
 #undef DEBUG_TYPE
 
   // 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());
-  }
+  if (targetOptions.getCompilationTarget() == gpu::CompilationTarget::Assembly)
+    return SmallVector<char, 0>(serializedISA->begin(), serializedISA->end());
 
   std::optional<SmallVector<char, 0>> result;
   moduleToObjectTimer.startTimer();

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,

diff  --git a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
index eabfd1c4d32eb..cae713a1ce1d2 100644
--- a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
@@ -130,6 +130,7 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {
 
     ASSERT_TRUE(
         StringRef(object->data(), object->size()).contains("nvvm_kernel"));
+    ASSERT_TRUE(StringRef(object->data(), object->size()).count('\0') == 0);
   }
 }