feat: emit labeled MLIR at each pipeline stage

solx-core/src/arguments.rs

@@ -160,7 +160,7 @@ pub struct Arguments {
     #[arg(long = "emit-llvm-ir", help_heading = "Output Selection")]
     pub output_llvm_ir: bool,
 
-    /// Emit MLIR (LLVM dialect, intermediate representation from Slang frontend).
+    /// Emit MLIR at each pipeline stage.
     /// Can be used with --output-dir to write .mlir files.
     #[cfg(feature = "mlir")]
     #[arg(long = "emit-mlir", help_heading = "Output Selection")]

solx-core/src/build/contract/mod.rs

@@ -5,6 +5,8 @@
 pub mod object;
 
 use std::collections::BTreeMap;
+#[cfg(feature = "mlir")]
+use std::collections::HashMap;
 use std::io::Write;
 use std::path::Path;
 use std::path::PathBuf;
@@ -42,9 +44,9 @@ pub struct Contract {
     pub legacy_assembly: Option<solx_evm_assembly::Assembly>,
     /// solc Yul IR.
     pub yul: Option<String>,
-    /// MLIR source code (LLVM dialect, from Slang frontend).
+    /// Labeled MLIR representations from each pipeline stage.
     #[cfg(feature = "mlir")]
-    pub mlir: Option<String>,
+    pub mlir: Option<HashMap<String, String>>,
 }
 
 impl Contract {
@@ -64,7 +66,7 @@ impl Contract {
         transient_storage_layout: Option<serde_json::Value>,
         legacy_assembly: Option<solx_evm_assembly::Assembly>,
         yul: Option<String>,
-        #[cfg(feature = "mlir")] mlir: Option<String>,
+        #[cfg(feature = "mlir")] mlir: Option<HashMap<String, String>>,
     ) -> Self {
         Self {
             name,
@@ -227,9 +229,11 @@ impl Contract {
             self.name.path.as_str(),
             self.name.name.as_deref(),
             solx_standard_json::InputSelector::MLIR,
-        ) && let Some(mlir) = self.mlir.take()
+        ) && let Some(stages) = self.mlir.take()
         {
-            writeln!(std::io::stdout(), "MLIR:\n{mlir}")?;
+            for (dialect, mlir_text) in stages.iter() {
+                writeln!(std::io::stdout(), "MLIR Dialect {dialect}:\n{mlir_text}")?;
+            }
         }
 
         if let Some(deploy_object_result) = self.deploy_object_result.as_mut()
@@ -788,16 +792,18 @@ impl Contract {
             self.name.path.as_str(),
             self.name.name.as_deref(),
             solx_standard_json::InputSelector::MLIR,
-        ) && let Some(mlir) = self.mlir.take()
+        ) && let Some(stages) = self.mlir.take()
         {
-            let output_name = format!(
-                "{contract_path}_{}.{}",
-                self.name.name.as_deref().unwrap_or(contract_name),
-                solx_utils::EXTENSION_MLIR,
-            );
-            let mut output_path = output_directory.to_owned();
-            output_path.push(output_name.as_str());
-            Self::write_to_file(output_path.as_path(), mlir, overwrite)?;
+            for (dialect, mlir_text) in stages.into_iter() {
+                let output_name = format!(
+                    "{contract_path}_{}.{dialect}.{}",
+                    self.name.name.as_deref().unwrap_or(contract_name),
+                    solx_utils::EXTENSION_MLIR,
+                );
+                let mut output_path = output_directory.to_owned();
+                output_path.push(output_name.as_str());
+                Self::write_to_file(output_path.as_path(), mlir_text, overwrite)?;
+            }
         }
         if let (Some(deploy_object_result), Some(runtime_object_result)) =
             (self.deploy_object_result, self.runtime_object_result)

solx-core/src/project/contract/mod.rs

@@ -7,6 +7,8 @@ pub mod metadata;
 
 use std::collections::BTreeMap;
 use std::collections::BTreeSet;
+#[cfg(feature = "mlir")]
+use std::collections::HashMap;
 
 #[cfg(feature = "mlir")]
 use anyhow::Context as _;
@@ -44,9 +46,9 @@ pub struct Contract {
     pub legacy_assembly: Option<solx_evm_assembly::Assembly>,
     /// solc Yul IR.
     pub yul: Option<String>,
-    /// solx MLIR IR.
+    /// Labeled MLIR representations from each pipeline stage.
     #[cfg(feature = "mlir")]
-    pub mlir: Option<String>,
+    pub mlir: Option<HashMap<String, String>>,
 }
 
 impl Contract {
@@ -65,7 +67,7 @@ impl Contract {
         transient_storage_layout: Option<serde_json::Value>,
         legacy_assembly: Option<solx_evm_assembly::Assembly>,
         yul: Option<String>,
-        #[cfg(feature = "mlir")] mlir: Option<String>,
+        #[cfg(feature = "mlir")] mlir: Option<HashMap<String, String>>,
     ) -> Self {
         Self {
             name,

solx-core/src/project/mod.rs

@@ -163,12 +163,21 @@ impl Project {
                     .and_then(|evm| evm.legacy_assembly.take());
 
                 #[cfg(feature = "mlir")]
-                let mlir_source = contract.mlir.clone();
+                let result = contract.mlir.as_ref().map(|stages| {
+                    stages
+                        .get(solx_mlir::Context::DIALECT_LLVM)
+                        .map(|source| {
+                            ContractIR::from(ContractMLIR {
+                                source: source.to_owned(),
+                            })
+                        })
+                        .ok_or_else(|| {
+                            anyhow::anyhow!("MLIR stages present but no LLVM dialect entry")
+                        })
+                        .map(Some)
+                });
                 #[cfg(feature = "mlir")]
-                let result = contract
-                    .mlir
-                    .take()
-                    .map(|source| Ok(Some(ContractIR::from(ContractMLIR { source }))));
+                let mlir_stages = contract.mlir.take();
                 #[cfg(not(feature = "mlir"))]
                 let result = if via_ir {
                     contract.ir.as_deref().map(|ir| {
@@ -209,7 +218,7 @@ impl Project {
                     legacy_assembly,
                     contract.ir,
                     #[cfg(feature = "mlir")]
-                    mlir_source,
+                    mlir_stages,
                 );
                 (name, Ok(contract))
             })

solx-mlir/src/context/mod.rs

@@ -45,6 +45,14 @@ pub struct Context<'context> {
 }
 
 impl<'context> Context<'context> {
+    // ---- Public constants ----
+
+    /// Dialect key for the Sol dialect MLIR stage.
+    pub const DIALECT_SOL: &'static str = "sol";
+
+    /// Dialect key for the LLVM dialect MLIR stage.
+    pub const DIALECT_LLVM: &'static str = "llvm";
+
     // ---- Private constants ----
 
     /// MLIR `builtin.module` operation name used to locate nested modules.
@@ -236,22 +244,30 @@ impl<'context> Context<'context> {
     }
 
     /// Consumes the context, runs the Sol-to-LLVM pass pipeline, and returns
-    /// the resulting LLVM dialect MLIR as text.
+    /// labeled MLIR representations captured at each pipeline stage.
+    ///
+    /// Returns `HashMap<dialect_name, mlir_text>` with entries for:
+    /// - `"sol"` — Full module in Sol dialect (only when `emit_mlir` is true)
+    /// - `"llvm"` — Runtime module only in LLVM dialect (always present)
     ///
     /// The Sol conversion pass produces a nested module structure:
     /// ```text
     /// module @Contract { deploy __entry + module @Contract_deployed { runtime __entry } }
     /// ```
     /// `solx-core` provides its own deploy wrapper, so this method extracts
-    /// only the inner `_deployed` module and renames it to
-    /// `runtime_code_identifier` so the LLVM module identifier matches what
-    /// `minimal_deploy_code` references.
+    /// only the inner module whose `sym_name` matches
+    /// `runtime_code_identifier` (the LLVM module identifier that
+    /// `minimal_deploy_code` references).
     ///
     /// # Errors
     ///
     /// Returns an error if re-parsing fails, the pass pipeline fails, or
     /// the deployed module is not found.
-    pub fn finalize_module(self, runtime_code_identifier: &str) -> anyhow::Result<String> {
+    pub fn finalize_module(
+        self,
+        runtime_code_identifier: &str,
+        emit_mlir: bool,
+    ) -> anyhow::Result<HashMap<String, String>> {
         // Re-parse the generated MLIR text to promote OperationBuilder
         // dictionary attributes to MLIR operation properties. Without this
         // round-trip, the Sol conversion pass fails because it expects
@@ -265,12 +281,22 @@ impl<'context> Context<'context> {
                 anyhow::anyhow!("failed to re-parse generated Sol dialect MLIR:\n{sol_text}")
             })?;
 
+        let mut stages = HashMap::new();
+
+        // Capture the Sol dialect MLIR before lowering (only when requested).
+        if emit_mlir {
+            stages.insert(
+                Self::DIALECT_SOL.to_owned(),
+                parsed_module.as_operation().to_string(),
+            );
+        }
+
         // Lower Sol → LLVM dialect.
         Self::run_sol_passes(self.builder.context, &mut parsed_module)?;
 
-        // Walk the outer module's body to find the inner `_deployed` module
-        // and extract it as the runtime code. Rename it so the LLVM module
-        // identifier matches what the deploy stub references.
+        // Walk the outer module's body to find the inner module whose
+        // `sym_name` matches `runtime_code_identifier` and extract it as
+        // the runtime code.
         let body = parsed_module.body();
         let mut deployed_operation = None;
         let mut operation = body.first_operation();
@@ -290,10 +316,15 @@ impl<'context> Context<'context> {
             operation = current.next_in_block();
         }
 
-        let runtime_op = deployed_operation
-            .ok_or_else(|| anyhow::anyhow!("no _deployed module in Sol pass output"))?;
+        let runtime_operation = deployed_operation.ok_or_else(|| {
+            anyhow::anyhow!(
+                "no module with sym_name `{runtime_code_identifier}` in Sol pass output"
+            )
+        })?;
+
+        stages.insert(Self::DIALECT_LLVM.to_owned(), runtime_operation.to_string());
 
-        Ok(runtime_op.to_string())
+        Ok(stages)
     }
 
     // ==== Phase 4: LLVM translation ====

solx-slang/src/slang/mod.rs

@@ -171,15 +171,20 @@ impl Frontend for Slang {
             };
 
             let runtime_code_identifier = format!("{contract_name}_deployed");
-            let mlir_source = context.finalize_module(&runtime_code_identifier)?;
+            let emit_mlir = input_json.settings.output_selection.check_selection(
+                file_identifier,
+                Some(contract_name.as_str()),
+                solx_standard_json::InputSelector::MLIR,
+            );
+            let mlir_stages = context.finalize_module(&runtime_code_identifier, emit_mlir)?;
 
             let evm = Some(solx_standard_json::output::contract::evm::EVM {
                 method_identifiers: Some(method_identifiers),
                 ..Default::default()
             });
 
             let contract = solx_standard_json::output::contract::Contract {
-                mlir: Some(mlir_source),
+                mlir: Some(mlir_stages),
                 evm,
                 ..Default::default()
             };

solx-standard-json/src/output/contract/mod.rs

@@ -4,6 +4,9 @@
 
 pub mod evm;
 
+#[cfg(feature = "mlir")]
+use std::collections::HashMap;
+
 use self::evm::EVM;
 
 ///
@@ -33,10 +36,10 @@ pub struct Contract {
     /// The contract Yul IR code.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub ir: Option<String>,
-    /// The contract MLIR source code.
+    /// Labeled MLIR representations from each pipeline stage.
     #[cfg(feature = "mlir")]
     #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub mlir: Option<String>,
+    pub mlir: Option<HashMap<String, String>>,
     /// The EVM data of the contract.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub evm: Option<EVM>,

solx/tests/cli/emit_mlir.rs

@@ -15,7 +15,12 @@ fn default() -> anyhow::Result<()> {
     ];
 
     let result = crate::cli::execute_solx(args)?;
-    result.success().stdout(predicate::str::contains("MLIR:"));
+    result
+        .success()
+        .stdout(predicate::str::contains("MLIR Dialect sol:"))
+        .stdout(predicate::str::contains("MLIR Dialect llvm:"))
+        .stdout(predicate::str::contains("sol.contract"))
+        .stdout(predicate::str::contains("llvm.func"));
 
     Ok(())
 }

solx/tests/cli/output_dir.rs

@@ -282,7 +282,23 @@ fn emit_mlir() -> anyhow::Result<()> {
         .filter(|entry| entry.path().extension().is_some_and(|ext| ext == "mlir"))
         .collect();
 
-    assert!(!entries.is_empty(), "Expected .mlir files to be created");
+    assert!(
+        entries.len() >= 2,
+        "Expected at least 2 .mlir files (one per dialect stage)"
+    );
+
+    let filenames: Vec<_> = entries
+        .iter()
+        .filter_map(|entry| entry.file_name().into_string().ok())
+        .collect();
+    assert!(
+        filenames.iter().any(|name| name.contains(".sol.mlir")),
+        "Expected a .sol.mlir file, found: {filenames:?}"
+    );
+    assert!(
+        filenames.iter().any(|name| name.contains(".llvm.mlir")),
+        "Expected a .llvm.mlir file, found: {filenames:?}"
+    );
 
     Ok(())
 }