feat(mlir): collect kernel metadata and surface via PTX comments

After makeLLIR runs the Triton allocation passes, read the MLIR module
attributes they write (ttg.shared, ttg.tensor_memory_size,
ttg.global_scratch_memory_*, ttg.profile_scratch_memory_*, and
ttg.total-num-warps) into a KernelMetadata struct on CudaBackend.
In makeASM, store the kernel entry-point name then append all fields as
structured `// meta:key=value` PTX line comments, which ptxas and the
CUDA driver ignore.

On the Rust side, add KernelMetadata with a parse() method that scans
those comment lines, and call it in compile_module() so each MlirModule
carries the parsed metadata alongside its PTX.

Author: arshadm

compiler/rustc_codegen_llvm/src/mlir/backend.rs

@@ -206,6 +206,18 @@ fn compile_module(mlir_module: &mut MlirModule<'static>) -> Result<(), MlirError
         .to_owned();
 
     eprintln!("Triton PTX output ({} bytes): {}", ptx.len(), ptx);
+
+    let metadata = crate::mlir::module::KernelMetadata::parse(&ptx);
+    eprintln!(
+        "Kernel metadata: name={:?} num_warps={} num_ctas={} shared={}B tmem={}B scratch={}B",
+        metadata.name,
+        metadata.num_warps,
+        metadata.num_ctas,
+        metadata.shared,
+        metadata.tmem_size,
+        metadata.global_scratch_size,
+    );
+    mlir_module.kernel_metadata = Some(metadata);
     mlir_module.ptx_asm = Some(ptx);
     Ok(())
 }

compiler/rustc_codegen_llvm/src/mlir/module.rs

@@ -25,6 +25,46 @@ use rustc_mlir::triton::TritonCompiler;
 
 use crate::mlir::backend::MlirCodegenBackend;
 
+/// Resource metadata for a compiled GPU kernel, recovered by parsing the
+/// structured `// meta:key=value` comments appended to the PTX by CudaBackend.
+#[derive(Debug, Default, Clone)]
+pub struct KernelMetadata {
+    pub name: String,
+    pub num_warps: i32,
+    pub num_ctas: i32,
+    pub shared: i32,
+    pub tmem_size: i32,
+    pub global_scratch_size: i32,
+    pub global_scratch_align: i32,
+    pub profile_scratch_size: i32,
+    pub profile_scratch_align: i32,
+}
+
+impl KernelMetadata {
+    /// Parse `// meta:key=value` lines from PTX output.
+    /// Lines that don't match the pattern are silently ignored.
+    pub fn parse(ptx: &str) -> Self {
+        let mut meta = KernelMetadata { num_ctas: 1, global_scratch_align: 1, profile_scratch_align: 1, ..Default::default() };
+        for line in ptx.lines() {
+            let Some(rest) = line.trim().strip_prefix("// meta:") else { continue };
+            let Some((key, val)) = rest.split_once('=') else { continue };
+            match key {
+                "name"                  => meta.name = val.to_owned(),
+                "num_warps"             => meta.num_warps            = val.parse().unwrap_or(0),
+                "num_ctas"              => meta.num_ctas             = val.parse().unwrap_or(1),
+                "shared"               => meta.shared               = val.parse().unwrap_or(0),
+                "tmem_size"             => meta.tmem_size            = val.parse().unwrap_or(0),
+                "global_scratch_size"   => meta.global_scratch_size  = val.parse().unwrap_or(0),
+                "global_scratch_align"  => meta.global_scratch_align = val.parse().unwrap_or(1),
+                "profile_scratch_size"  => meta.profile_scratch_size = val.parse().unwrap_or(0),
+                "profile_scratch_align" => meta.profile_scratch_align= val.parse().unwrap_or(1),
+                _ => {}
+            }
+        }
+        meta
+    }
+}
+
 /// Represents an MLIR module during codegen
 pub struct MlirModule<'c> {
     pub name: String,
@@ -36,6 +76,8 @@ pub struct MlirModule<'c> {
     pub ptx_asm: Option<String>,
     /// MLIR source captured after cleanup passes, before Triton passes run.
     pub mlir_source: Option<String>,
+    /// Kernel metadata parsed from the PTX comment block appended by CudaBackend.
+    pub kernel_metadata: Option<KernelMetadata>,
 }
 
 unsafe impl<'c> Send for MlirModule<'c> {}
@@ -57,7 +99,7 @@ impl<'c> MlirModule<'c> {
         let compiler = TritonCompiler::new(context.to_raw(), "cuda", &options)
             .expect("Failed to create Triton compiler");
 
-        Self { name: mod_name.to_string(), mlir: module, compiler, context, ptx_asm: None, mlir_source: None }
+        Self { name: mod_name.to_string(), mlir: module, compiler, context, ptx_asm: None, mlir_source: None, kernel_metadata: None }
     }
 
     pub fn context(&self) -> &Context {
@@ -84,6 +126,7 @@ impl<'c> MlirModule<'c> {
             compiler,
             ptx_asm: None,
             mlir_source: None,
+            kernel_metadata: None,
         }
     }
 

compiler/rustc_mlir/mlir-wrapper/triton/backend/CudaBackend.cpp

@@ -219,6 +219,7 @@ LogicalResult CudaBackend::makeASM(MLIRContext &context, ModuleOp module) {
     llvm::errs() << "Could not find kernel name in PTX output\n";
     return LogicalResult::failure();
   }
+  m_metadata.name = kernel_name;
 
   // 4. Post-process version and target
   char ptx_major_minor[8];
@@ -233,6 +234,20 @@ LogicalResult CudaBackend::makeASM(MLIRContext &context, ModuleOp module) {
   // No 'knobs' defined; always remove for now or leave as TODO.
   ret = std::regex_replace(ret, std::regex(R"(,\s*debug|debug,\s*)"), "");
 
+  // 6. Append kernel metadata as PTX line comments. PTX comments are ignored
+  //    by ptxas and the CUDA driver, so this is safe. The Rust side parses
+  //    these lines to recover launch parameters without a separate FFI channel.
+  ret += "\n// --- triton-metadata ---\n";
+  ret += "// meta:name="                  + m_metadata.name + "\n";
+  ret += "// meta:num_warps="             + std::to_string(m_metadata.num_warps) + "\n";
+  ret += "// meta:num_ctas="              + std::to_string(m_metadata.num_ctas) + "\n";
+  ret += "// meta:shared="               + std::to_string(m_metadata.shared) + "\n";
+  ret += "// meta:tmem_size="             + std::to_string(m_metadata.tmem_size) + "\n";
+  ret += "// meta:global_scratch_size="   + std::to_string(m_metadata.global_scratch_size) + "\n";
+  ret += "// meta:global_scratch_align="  + std::to_string(m_metadata.global_scratch_align) + "\n";
+  ret += "// meta:profile_scratch_size="  + std::to_string(m_metadata.profile_scratch_size) + "\n";
+  ret += "// meta:profile_scratch_align=" + std::to_string(m_metadata.profile_scratch_align) + "\n";
+
   // 7. Save PTX (exposed via getASM())
   m_asm = std::move(ret);
   return LogicalResult::success();
@@ -474,45 +489,30 @@ LogicalResult CudaBackend::makeLLIR(MLIRContext &context, ModuleOp module) {
     // CUDABackend.instrumentation.patch("llvmir_to_llvm", pm, mod.context)
   }
 
-  return pm.run(op);
+  auto result = pm.run(op);
+  if (succeeded(result)) {
+    // Read resource metadata written by the allocation passes as MLIR module
+    // attributes. These are later appended to the PTX as comments so Rust can
+    // recover them without a separate FFI channel.
+    auto getInt = [&](llvm::StringRef key, int32_t def = 0) -> int32_t {
+      auto attr = op->getAttrOfType<mlir::IntegerAttr>(key);
+      return attr ? static_cast<int32_t>(attr.getInt()) : def;
+    };
 
-  //  # LLVM-IR (MLIR) -> LLVM-IR (LLVM)
-  //       llvm.init_targets()
-  //       context = llvm.context()
-  //       if knobs.compilation.enable_asan:
-  //           raise RuntimeError(
-  //               "Address Sanitizer Error: Address sanitizer is currently only
-  //               supported on the AMD backend")
-  //       llvm_mod = llvm.to_module(mod, context)
-  //       proc = sm_arch_from_capability(capability)
-  //       features = get_features(options, self.target.arch)
-  //       triple = 'nvptx64-nvidia-cuda'
-  //       nvidia.set_short_ptr()
-  //       llvm.attach_datalayout(llvm_mod, triple, proc, features)
-  //       nvidia.set_nvvm_reflect_ftz(llvm_mod)
-
-  //       if options.extern_libs and nvidia.has_extern_deps(llvm_mod):
-  //           paths = [path for (name, path) in options.extern_libs]
-  //           llvm.link_extern_libs(llvm_mod, paths)
-
-  //       llvm.optimize_module(llvm_mod, llvm.OPTIMIZE_O3)
-
-  //       # Get some metadata
-  //       # warp-specialization mutates num_warps
-  //       total_num_warps = src.get_int_attr("ttg.total-num-warps")
-  //       if total_num_warps is not None:
-  //           metadata["num_warps"] = total_num_warps
-  //       metadata["shared"] = src.get_int_attr("ttg.shared")
-  //       metadata["tmem_size"] = src.get_int_attr("ttg.tensor_memory_size")
-  //       metadata["global_scratch_size"] =
-  //       src.get_int_attr("ttg.global_scratch_memory_size")
-  //       metadata["global_scratch_align"] =
-  //       src.get_int_attr("ttg.global_scratch_memory_alignment")
-  //       metadata["profile_scratch_size"] =
-  //       src.get_int_attr("ttg.profile_scratch_memory_size") or 0
-  //       metadata["profile_scratch_align"] =
-  //       src.get_int_attr("ttg.profile_scratch_memory_alignment") or 1 ret =
-  //       str(llvm_mod) del llvm_mod del context return ret
+    // Warp-specialization may have mutated the warp count, so prefer the
+    // post-pipeline attribute over the original options value.
+    auto totalWarps = op->getAttrOfType<mlir::IntegerAttr>("ttg.total-num-warps");
+    m_metadata.num_warps   = totalWarps ? static_cast<int32_t>(totalWarps.getInt())
+                                        : m_options.num_warps;
+    m_metadata.num_ctas              = m_options.num_ctas;
+    m_metadata.shared                = getInt("ttg.shared");
+    m_metadata.tmem_size             = getInt("ttg.tensor_memory_size");
+    m_metadata.global_scratch_size   = getInt("ttg.global_scratch_memory_size");
+    m_metadata.global_scratch_align  = getInt("ttg.global_scratch_memory_alignment", 1);
+    m_metadata.profile_scratch_size  = getInt("ttg.profile_scratch_memory_size");
+    m_metadata.profile_scratch_align = getInt("ttg.profile_scratch_memory_alignment", 1);
+  }
+  return result;
 }
 
 std::unique_ptr<mlir::Pass>

compiler/rustc_mlir/mlir-wrapper/triton/backend/CudaBackend.h

@@ -37,6 +37,26 @@ namespace triton {
 
 namespace ttng = mlir::triton::nvidia_gpu;
 
+// ---------------------------------------------------------------------------
+// Kernel metadata collected during lowering
+// ---------------------------------------------------------------------------
+
+/// Resource metadata for a compiled GPU kernel. Populated by makeLLIR (from
+/// MLIR module attributes written by the allocation passes) and makeASM (kernel
+/// name extracted from PTX). Appended to m_asm as PTX line comments so that
+/// the Rust side can parse them without an extra FFI channel.
+struct KernelMetadata {
+  std::string name;
+  int32_t num_warps            = 0;
+  int32_t num_ctas             = 1;
+  int32_t shared               = 0;  // shared memory in bytes
+  int32_t tmem_size            = 0;  // tensor memory (Blackwell SM100+)
+  int32_t global_scratch_size  = 0;
+  int32_t global_scratch_align = 1;
+  int32_t profile_scratch_size  = 0;
+  int32_t profile_scratch_align = 1;
+};
+
 // ---------------------------------------------------------------------------
 // FFI-safe helper types
 // These map directly to #[repr(C)] Rust structs with equivalent fields.
@@ -216,6 +236,7 @@ class CudaBackend : public Backend {
 
   CudaCompileOptions m_options;
   Capability m_capability;
+  KernelMetadata m_metadata;
 
   std::unordered_map<CudaPass, std::unique_ptr<Pass> (*)()> m_nvidia_pass_fns =
       {