runtime: validate compute yield context and keep int literals inferred

Author: darmie

crates/compiler/src/ssa.rs

@@ -69,6 +69,9 @@ pub struct SsaBuilder {
     /// External function link names (alias -> ZRTL symbol)
     /// e.g., "tensor_add" -> "$Tensor$add"
     extern_link_names: IndexMap<InternedString, String>,
+    /// Captured yield values for active compute-expression lowering contexts.
+    /// Empty outside compute expression translation.
+    compute_yield_stack: Vec<Vec<HirId>>,
 }
 
 /// Context for pattern matching
@@ -314,6 +317,7 @@ impl SsaBuilder {
             address_taken_vars: HashSet::new(),
             stack_slots: IndexMap::new(),
             extern_link_names: IndexMap::new(),
+            compute_yield_stack: Vec::new(),
         }
     }
 
@@ -1368,8 +1372,16 @@ impl SsaBuilder {
                 self.translate_expression(block_id, expr)?;
             }
             TypedStatement::Yield(expr) => {
-                // Yield behaves like an expression statement at the SSA level.
-                self.translate_expression(block_id, expr)?;
+                if self.compute_yield_stack.is_empty() {
+                    return Err(crate::CompilerError::Analysis(
+                        "`yield` is only valid inside compute expression bodies".to_string(),
+                    ));
+                }
+
+                let yielded = self.translate_expression(block_id, expr)?;
+                if let Some(active_yields) = self.compute_yield_stack.last_mut() {
+                    active_yields.push(yielded);
+                }
             }
 
             TypedStatement::Match(match_stmt) => {
@@ -1494,6 +1506,31 @@ impl SsaBuilder {
         Ok(())
     }
 
+    fn translate_compute_dispatch_call(
+        &mut self,
+        block_id: HirId,
+        expr: &TypedNode<TypedExpression>,
+        compute: &zyntax_typed_ast::typed_ast::TypedComputeExpr,
+    ) -> CompilerResult<HirId> {
+        let compute_name = InternedString::new_global(INTERNAL_COMPUTE_ALIAS);
+        let lowered = zyntax_typed_ast::typed_ast::TypedCall {
+            callee: Box::new(zyntax_typed_ast::typed_node(
+                TypedExpression::Variable(compute_name),
+                Type::Any,
+                expr.span,
+            )),
+            positional_args: compute.args.clone(),
+            named_args: vec![],
+            type_args: vec![],
+        };
+        let lowered_expr = zyntax_typed_ast::typed_node(
+            TypedExpression::Call(lowered),
+            expr.ty.clone(),
+            expr.span,
+        );
+        self.translate_expression(block_id, &lowered_expr)
+    }
+
     /// Translate expression to SSA value
     fn translate_expression(
         &mut self,
@@ -1941,25 +1978,49 @@ impl SsaBuilder {
             }
 
             TypedExpression::Compute(compute) => {
-                // Lower compute expressions through the existing call pipeline for now.
-                // This preserves call-based execution while keeping typed compute structure.
-                let compute_name = InternedString::new_global(INTERNAL_COMPUTE_ALIAS);
-                let lowered = zyntax_typed_ast::typed_ast::TypedCall {
-                    callee: Box::new(zyntax_typed_ast::typed_node(
-                        TypedExpression::Variable(compute_name),
-                        Type::Any,
-                        expr.span,
-                    )),
-                    positional_args: compute.args.clone(),
-                    named_args: vec![],
-                    type_args: vec![],
-                };
-                let lowered_expr = zyntax_typed_ast::typed_node(
-                    TypedExpression::Call(lowered),
-                    expr.ty.clone(),
-                    expr.span,
-                );
-                self.translate_expression(block_id, &lowered_expr)
+                let has_direct_yield = compute.body.statements.iter().any(|stmt| {
+                    matches!(
+                        stmt.node,
+                        zyntax_typed_ast::typed_ast::TypedStatement::Yield(_)
+                    )
+                });
+
+                if !has_direct_yield {
+                    // Preserve legacy runtime behavior for non-reduction compute blocks.
+                    return self.translate_compute_dispatch_call(block_id, expr, compute);
+                }
+
+                // CPU fallback path for reduction-style compute blocks.
+                // Evaluate explicit compute args for side effects, then execute compute body.
+                for arg in &compute.args {
+                    self.translate_expression(block_id, arg)?;
+                }
+
+                self.compute_yield_stack.push(Vec::new());
+                let mut current_block = block_id;
+                for stmt in &compute.body.statements {
+                    match &stmt.node {
+                        zyntax_typed_ast::typed_ast::TypedStatement::Let(_)
+                        | zyntax_typed_ast::typed_ast::TypedStatement::Expression(_)
+                        | zyntax_typed_ast::typed_ast::TypedStatement::Yield(_) => {
+                            current_block = self.process_statement(current_block, stmt)?;
+                        }
+                        _ => {
+                            self.compute_yield_stack.pop();
+                            return Err(crate::CompilerError::Analysis(
+                                "compute body fallback currently supports only let/expression/yield statements"
+                                    .to_string(),
+                            ));
+                        }
+                    }
+                }
+
+                let yields = self.compute_yield_stack.pop().unwrap_or_default();
+                if let Some(last_yield) = yields.last().copied() {
+                    Ok(last_yield)
+                } else {
+                    self.translate_compute_dispatch_call(block_id, expr, compute)
+                }
             }
 
             TypedExpression::Field(field_access) => {

crates/zyn_peg/src/runtime2/interpreter.rs

@@ -809,11 +809,11 @@ impl<'g> GrammarInterpreter<'g> {
             _ => return Err(format!("unknown TypedExpression variant: {}", variant)),
         };
 
-        // Determine the type based on the expression variant
-        // Default to I64 for integers (64-bit system) and F64 for floats
+        // Determine the type based on the expression variant.
+        // Integer literals default to I32; assignment/context-driven inference can refine later.
         let ty = match &expr {
             TypedExpression::Literal(TypedLiteral::Integer(_)) => {
-                Type::Primitive(PrimitiveType::I64)
+                Type::Primitive(PrimitiveType::I32)
             }
             TypedExpression::Literal(TypedLiteral::Float(_)) => Type::Primitive(PrimitiveType::F64),
             TypedExpression::Literal(TypedLiteral::String(_)) => {
@@ -824,6 +824,7 @@ impl<'g> GrammarInterpreter<'g> {
             // Use Type::Any to signal that lowering should infer the type
             TypedExpression::Call(_) => Type::Any,
             TypedExpression::Variable(_) => Type::Any,
+            TypedExpression::Compute(_) => Type::Any,
             // Struct literal gets its type from the struct name - use Unresolved for compiler to resolve
             TypedExpression::Struct(lit) => Type::Unresolved(lit.name),
             _ => Type::Primitive(PrimitiveType::Unit),

crates/zynml/tests/e2e_tests.rs

@@ -3590,6 +3590,55 @@ mod execution {
         }
     }
 
+    #[test]
+    fn test_execute_compute_yield_fallback() {
+        let Some(mut zynml) = create_runtime_with_plugins() else {
+            println!("Skipping: plugins not available");
+            return;
+        };
+
+        let source = r#"
+            fn compute_fallback() {
+                let x = 41
+                let _y = compute(x) @device("cpu") {
+                    yield x
+                }
+            }
+        "#;
+
+        let functions = zynml
+            .load_source(source)
+            .expect("compute fallback program should compile");
+        assert!(
+            functions.iter().any(|f| f == "compute_fallback"),
+            "compute_fallback should be exported"
+        );
+    }
+
+    #[test]
+    fn test_execute_yield_outside_compute_rejected() {
+        let Some(mut zynml) = create_runtime_with_plugins() else {
+            println!("Skipping: plugins not available");
+            return;
+        };
+
+        let source = r#"
+            fn invalid_yield() {
+                yield 1
+            }
+        "#;
+
+        let functions = zynml
+            .load_source(source)
+            .expect("compiler should continue while skipping invalid function");
+        let message = format!("{:?}", functions);
+        assert!(
+            functions.is_empty(),
+            "yield outside compute should prevent function export, got: {}",
+            message
+        );
+    }
+
     // Regression coverage for the full hello example (prelude + tensor + dynamic println).
     // Keep panic isolation via catch_unwind so runtime regressions produce actionable test output.
     #[test]