[Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM) #28595

andportnoy · 2025-05-07T21:52:34Z

Separate commits to make it easier to review.

apaszke · 2025-05-12T06:09:27Z

jaxlib/mosaic/gpu/custom_call.cc

+    }
+    stdout += buf;
+  }
+  if (close(stdout_pipe[0]) == -1) {


Please use absl::Cleanup to close the pipes. This, or the close above might not get executed if something else fails along the way

apaszke · 2025-05-12T06:10:36Z

jaxlib/mosaic/gpu/custom_call.cc

+      return absl::InternalError(
+          absl::StrCat("Failed to read from pipe: ", strerror(errno)));
+    }
+    stdout += buf;


This looks sketchy to me? The addition will probably look for the null byte to decide where to stop the concatenation but there might be no null byte in there (read does not insert one). Use stdout.append(buf, status) (and perhaps rename status to bytes_read)

Agreed. I'm more used to writing C for systems code like that and forming the string explicitly.

apaszke · 2025-05-12T06:12:11Z

jaxlib/mosaic/gpu/custom_call.cc

    return absl::InternalError("Failed to wait for CUDA tool invocation");
  }
-  if (status != 0) return absl::InternalError("CUDA tool failed");
+  if (status != 0) return absl::InternalError(stdout);


Please prepend a failure message to stdout in case the tool didn't write something helpful

apaszke · 2025-05-12T06:13:53Z

jaxlib/mosaic/gpu/target.cc


-  const std::string sm = sm_arch_specific ? sm_arch_specific : sm_base;
-
+absl::StatusOr<std::string> GetPtxIsaVersion(int major, int minor) {


This is dead now, isn't it? Perhaps delete it?

Sorry this PR is a bit raw, I'll refactor it a bit better + add more inline docs. This function does play a role here, it takes an input PTX ISA version specified as major.minor and returns min(major.minor, <latest PTX ISA supported by LLVM>). This is to cover cases when ptxas is more recent and LLVM hasn't caught up to it yet.

I refactored the function for clarity.

jaxlib/mosaic/gpu/custom_call.cc

jreiffers · 2025-05-23T06:30:17Z

jaxlib/mosaic/gpu/custom_call.cc

+absl::StatusOr<int> GetLatestPtxasPtxIsaVersion() {
+  std::vector<const char*> ptxas_args = {"ptxas", "--input-as-string",
+                                         ".version 99.99", nullptr};
+  auto result = RunCUDATool("ptxas", ptxas_args);


Nit: Avoid using auto when the type isn't obvious. Something like

auto status = RunCUDATool("ptxas", ptxas_args).status();

Might be fine.

Thank you, fixed.

jreiffers · 2025-05-23T06:35:00Z

jaxlib/mosaic/gpu/custom_call.cc

+        absl::StrFormat("Failed to parse PTX ISA minor version, expected a "
+                        "parsable integer, instead got: %s",
+                        major_minor[1]));
+  }


We might want to fail loudly here if minor is ever >= 10.

jreiffers · 2025-05-23T06:36:08Z

jaxlib/mosaic/gpu/custom_call.cc

+
+absl::StatusOr<std::string> GetPtxIsaVersion() {
+  int ptxas_latest_version;
+  TF_ASSIGN_OR_RETURN(ptxas_latest_version, GetLatestPtxasPtxIsaVersion());


Nit: why not just put the declaration in here?

TF_ASSIGN_OR_RETURN(int ptxas_latest_version, GetLatestPtxasPtxIsaVersion());

Fixed, thank you.

jreiffers · 2025-05-23T06:42:23Z

jaxlib/mosaic/gpu/target.cc

  for (const llvm::SubtargetFeatureKV& feature :
-       subtarget_info->getEnabledProcessorFeatures()) {
+       subtarget_info->getAllProcessorFeatures()) {
    if (absl::StartsWith(feature.Key, "ptx")) {


Nothing really wrong here, but it's not super idiomatic:

absl::string_view key = feature.Key; if (absl::ConsumePrefix(&key, "ptx")) { // No need for version_string }

Cool, didn't know about ConsumePrefix.

jaxlib/mosaic/gpu/custom_call.cc

apaszke · 2025-05-26T08:58:13Z

jaxlib/mosaic/gpu/custom_call.cc

      continue;
    }
+    // Dump SASS.
+    std::cout << *result << std::endl;


I think you missed an explicit print for ptxas output above (if dump_ptxas is on)

Fixed that (see last force push) and squashed the commits.

andportnoy added the CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits label May 7, 2025

andportnoy changed the title ~~Mosaic gpu ptx isa from ptxas and llvm~~ [Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM) May 7, 2025

andportnoy requested a review from apaszke May 7, 2025 22:44

apaszke reviewed May 12, 2025

View reviewed changes

Amir-19 mentioned this pull request May 12, 2025

[mosaic-gpu] add multicast ptr support to TMA with overlapped gemm and all reduce examples #28679

Open

apaszke self-assigned this May 22, 2025

andportnoy force-pushed the mosaic-gpu-ptx-isa-from-ptxas-and-llvm branch from 8413c48 to ff213dc Compare May 22, 2025 21:29

andportnoy requested a review from apaszke May 22, 2025 23:10

jreiffers reviewed May 23, 2025

View reviewed changes

jaxlib/mosaic/gpu/custom_call.cc Show resolved Hide resolved

jreiffers reviewed May 23, 2025

View reviewed changes

apaszke approved these changes May 26, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels May 26, 2025

[Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM)

c1e8f25

andportnoy force-pushed the mosaic-gpu-ptx-isa-from-ptxas-and-llvm branch from b8d2305 to c1e8f25 Compare May 26, 2025 15:18

copybara-service bot merged commit 2cbec58 into jax-ml:main May 27, 2025
20 of 23 checks passed


		const std::string sm = sm_arch_specific ? sm_arch_specific : sm_base;

		absl::StatusOr<std::string> GetPtxIsaVersion(int major, int minor) {

[Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM) #28595

[Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM) #28595

Uh oh!

Conversation

andportnoy commented May 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jreiffers May 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

andportnoy commented May 7, 2025 •

edited

Loading

jreiffers May 23, 2025 •

edited

Loading