Fixed incorrect PTX parsing of ret instruction after branch label (#…

…17859)

The PTX parser replaces PTX code with inline PTX code (using inline ASM blocks).
It considers a branch label and the immediate instruction as a single unit to process.  
During the ASM->CUDA transform step,  it searches for the `ret` instruction in the string and replaces the whole statement and not the substring that contains the `ret;` instruction. which means an expression like:

```asm

BB0_1:
ret;
```

is parsed as: 

```asm 

BB0_1: ret;

```

and then transformed to: 

```asm

bra RETTGT

``` 

instead of:

```asm 

BB0_1: bra RETTGT

```

This merge request fixes this bug.

Authors:
  - Basit Ayantunde (https://github.com/lamarrr)

Approvers:
  - David Wendt (https://github.com/davidwendt)
  - Shruti Shivakumar (https://github.com/shrshi)

URL: #17859

cpp/src/jit/parser.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -375,9 +375,14 @@ std::string ptx_parser::parse()
 
   // Don't use std::accumulate until C++20 when rvalue references are supported
   auto final_output = fn_header_output + "\n asm volatile (\"{\");";
-  for (auto const& line : fn_body_output)
-    final_output += line.find("ret;") != std::string::npos ? "  asm volatile (\"bra RETTGT;\");\n"
-                                                           : "  " + line + "\n";
+  for (auto const& line : fn_body_output) {
+    std::string output{line};
+    std::string_view const ret_instruction = "ret;";
+    if (auto start = output.find(ret_instruction); start != std::string::npos) {
+      output.replace(start, ret_instruction.size(), "bra RETTGT;");
+    }
+    final_output += "  " + output + "\n";
+  }
   return final_output + " asm volatile (\"RETTGT:}\");}";
 }
 

cpp/tests/jit/parse_ptx_function.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -69,6 +69,7 @@ __device__ __inline__ void GENERIC_OP(
 ){
  asm volatile ("{");
  asm volatile ("bra RETTGT;");
+ /** ret*/
  asm volatile ("RETTGT:}");}
 )";
 
@@ -78,6 +79,28 @@ __device__ __inline__ void GENERIC_OP(
   EXPECT_TRUE(ptx_equal(cuda_source, expected));
 }
 
+TEST_F(JitParseTest, PTXWithBranchLabel)
+{
+  std::string raw_ptx = R"(
+.visible .func _Z1flPaS_(){
+BB0:
+  ret;
+}
+)";
+
+  std::string expected = R"(
+__device__ __inline__ void GENERIC_OP(){
+ asm volatile ("{");
+ asm volatile ("BB0: bra RETTGT;");
+ /** BB0: ret*/
+ asm volatile ("RETTGT:}");}
+)";
+
+  std::string cuda_source = cudf::jit::parse_single_function_ptx(raw_ptx, "GENERIC_OP", {});
+
+  EXPECT_TRUE(ptx_equal(cuda_source, expected));
+}
+
 TEST_F(JitParseTest, PTXWithPragma)
 {
   std::string raw_ptx = R"(