[Mosaic GPU] Improve default kernel name and add option to customize

This allows users to distinguish Mosaic GPU kernels from other kernels
when using profiling programs such as Nsight Systems.

The new default behavior is to use `mosaic_gpu_<def_name>_kernel` as
the kernel name, where `<def_name>` is the name of the Mosaic GPU
Python kernel function passed to `as_gpu_kernel` or
`as_torch_gpu_kernel`.

We also add a new `kernel_name` optional argument to `as_gpu_kernel`
and `as_torch_gpu_kernel`. If `kernel_name` is not `None`, the
resulting kernel name is `mosaic_gpu_<kernel_name>_kernel`. This is
useful when the Mosaic GPU Python kernel function is constructed
through metaprogramming so that the final specialized kernel can have
different meaningful names depending on the metaparameters.

Previously the kernel name was always `main_kernel`.

Author: andportnoy

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