Merge pull request jax-ml#25006 from andportnoy:aportnoy/mosaic-gpu-kernel-custom-name

PiperOrigin-RevId: 702772768

Author: Google-ML-Automation

jax/experimental/mosaic/gpu/core.py

@@ -847,6 +847,7 @@ def _lower_as_gpu_kernel(
     out_shape,
     smem_scratch_shape: ShapeTree | Union[ShapeTree],
     module_name: str,
+    kernel_name: str | None = None,
     prof_spec: profiler.ProfilerSpec | None = None,
 ):
   ptr_ty = ir.Type.parse("!llvm.ptr")
@@ -873,6 +874,8 @@ def _shape_to_ref_ty(shape: jax.ShapeDtypeStruct) -> ir.MemRefType:
   module = ir.Module.create()
   attrs = module.operation.attributes
   attrs["sym_name"] = ir.StringAttr.get(module_name)
+  if kernel_name is None:
+    kernel_name = getattr(body, "__name__", "anonymous")
   with ir.InsertionPoint(module.body):
     _declare_runtime_functions()
     gmem_scratch_bytes = 0
@@ -882,7 +885,7 @@ def _shape_to_ref_ty(shape: jax.ShapeDtypeStruct) -> ir.MemRefType:
         ir.Attribute.parse("#llvm.linkage<external>"),
         addr_space=ir.IntegerAttr.get(i32, 4),  # GPU constant memory.
     )
-    @func.FuncOp.from_py_func(ptr_ty, ptr_ty)
+    @func.FuncOp.from_py_func(ptr_ty, ptr_ty, name=f"mosaic_gpu_{kernel_name}")
     def main(token_ptr, buffers):
       nonlocal gmem_scratch_bytes
       token = builtin.unrealized_conversion_cast([token_ty], [token_ptr])
@@ -947,6 +950,7 @@ def as_gpu_kernel(
     prof_spec: profiler.ProfilerSpec | None = None,
     cluster: tuple[int, int, int] = (1, 1, 1),
     module_name: str = "unknown",
+    kernel_name: str | None = None,
 ):
   if isinstance(in_shape, list):
     in_shape = tuple(in_shape)
@@ -956,7 +960,7 @@ def as_gpu_kernel(
   module, out_shape, unwrap_output_tuple = (
       _lower_as_gpu_kernel(
           body, grid, cluster, block, in_shape, out_shape, smem_scratch_shape,
-          module_name, prof_spec
+          module_name, kernel_name, prof_spec
       )
   )
 
@@ -1014,6 +1018,7 @@ def as_torch_gpu_kernel(
     prof_spec: profiler.ProfilerSpec | None = None,
     cluster: tuple[int, int, int] = (1, 1, 1),
     module_name: str = "unknown",
+    kernel_name: str | None = None,
 ):
   try:
     import torch
@@ -1032,7 +1037,7 @@ def as_torch_gpu_kernel(
   module, out_shape, unwrap_output_tuple = (
       _lower_as_gpu_kernel(
           body, grid, cluster, block, in_shape, out_shape, smem_scratch_shape,
-          module_name, prof_spec
+          module_name, kernel_name, prof_spec
       )
   )
 

jaxlib/mosaic/gpu/BUILD

@@ -126,6 +126,7 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/synchronization",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",

jaxlib/mosaic/gpu/custom_call.cc

@@ -38,6 +38,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
 #include "absl/synchronization/mutex.h"
 #include "llvm/include/llvm/ADT/SmallVector.h"
 #include "llvm/include/llvm/Support/CodeGen.h"
@@ -415,6 +416,40 @@ GetKernelCache() {
   return std::make_pair(&context_cache, &mutex);
 }
 
+absl::StatusOr<std::pair<std::string, std::string>> GetHostAndInitFuncNames(
+    mlir::ModuleOp module_op) {
+  // We look for two top level C-interface functions:
+  // - "host" function with symbol name "_mlir_ciface_<foo>"
+  // - "init" function with symbol name "_mlir_ciface_<foo>_init"
+  constexpr std::string_view prefix = "_mlir_ciface_";
+  std::vector<std::string> names;
+  for (mlir::LLVM::LLVMFuncOp llvm_func :
+       module_op.getOps<mlir::LLVM::LLVMFuncOp>()) {
+    if (llvm_func.getName().starts_with(prefix)) {
+      names.push_back(llvm_func.getName().str());
+    }
+  }
+  if (auto size = names.size(); size != 2) {
+    return absl::InternalError(absl::StrFormat(
+        "Expected to locate 2 symbols with %s prefix in the MLIR module, found "
+        "%d instead.",
+        prefix, size));
+  }
+  // _mlir_ciface_<foo>_init now follows _mlir_ciface_<foo>
+  std::sort(names.begin(), names.end());
+
+  std::string host_func_name = names[0];
+  std::string init_func_name = names[1];
+
+  if (init_func_name != absl::StrCat(host_func_name, "_init")) {
+    return absl::InternalError(absl::StrFormat(
+        "Expected init function name to equal the concatenation of the host "
+        "function name and the \"_init\" suffix, instead got "
+        "init_func_name=%s, host_func_name=%s.",
+        init_func_name, host_func_name));
+  }
+  return std::make_pair(host_func_name, init_func_name);
+}
 
 absl::StatusOr<CompiledKernel> CompileAndInit(const char* module) {
   mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
@@ -430,9 +465,16 @@ absl::StatusOr<CompiledKernel> CompileAndInit(const char* module) {
     return maybe_engine.status();
   }
   mlir::ExecutionEngine* execution_engine = maybe_engine->get();
-  auto main = execution_engine->lookupPacked("_mlir_ciface_main");
-  auto init = execution_engine->lookupPacked("_mlir_ciface_main_init");
-  if (!init || !main) {
+
+  auto host_and_init_func_names = GetHostAndInitFuncNames(*module_op);
+  if (!host_and_init_func_names.ok()) {
+    return host_and_init_func_names.status();
+  }
+  auto [host_name, init_name] = host_and_init_func_names.value();
+
+  auto host = execution_engine->lookupPacked(host_name);
+  auto init = execution_engine->lookupPacked(init_name);
+  if (!init || !host) {
     return absl::InternalError("Failed to retrieve kernel function");
   }
   void* module_ptr = nullptr;
@@ -442,7 +484,7 @@ absl::StatusOr<CompiledKernel> CompileAndInit(const char* module) {
   void*** init_args[2] = {&module_ptr_ptr, &kernel_ptr_ptr};
   reinterpret_cast<MosaicInitFunc*>(*init)(init_args);
   return CompiledKernel(std::move(*maybe_engine), kernel_ptr,
-                          reinterpret_cast<MosaicHostFunc*>(*main));
+                        reinterpret_cast<MosaicHostFunc*>(*host));
 }
 
 // Each compiled kernel has a unique init func, and each kernel is used from