From 47426df76a7adfb2857cb8c20bed8d32d4842333 Mon Sep 17 00:00:00 2001 From: David Hagen Date: Tue, 11 Aug 2026 17:34:21 -0400 Subject: [PATCH] Fix bug in nested LLVM And and Or nodes --- src/tensora/codegen/_ir_to_llvm.py | 22 ++++++++++--------- tests/test_combinatorically.py | 34 ++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 10 deletions(-) diff --git a/src/tensora/codegen/_ir_to_llvm.py b/src/tensora/codegen/_ir_to_llvm.py index 033a370..817cc28 100644 --- a/src/tensora/codegen/_ir_to_llvm.py +++ b/src/tensora/codegen/_ir_to_llvm.py @@ -289,44 +289,46 @@ def ir_to_llvm_less_than_or_equal( def ir_to_llvm_and( self: And, builder: llvm.IRBuilder, locals: dict[str, llvm.Value] ) -> llvm.Value: - left_block = builder.block + left = ir_to_llvm_expression(self.left, builder, locals) + left_end_block = builder.block + right_block = builder.append_basic_block() end_block = builder.append_basic_block() - left = ir_to_llvm_expression(self.left, builder, locals) - builder.cbranch(left, right_block, end_block) builder.position_at_end(right_block) right = ir_to_llvm_expression(self.right, builder, locals) + right_end_block = builder.block builder.branch(end_block) builder.position_at_end(end_block) phi = builder.phi(llvm_boolean_type) - phi.add_incoming(llvm.Constant(llvm_boolean_type, 0), left_block) - phi.add_incoming(right, right_block) + phi.add_incoming(llvm.Constant(llvm_boolean_type, 0), left_end_block) + phi.add_incoming(right, right_end_block) return phi @ir_to_llvm_expression.register(Or) def ir_to_llvm_or(self: Or, builder: llvm.IRBuilder, locals: dict[str, llvm.Value]) -> llvm.Value: - left_block = builder.block + left = ir_to_llvm_expression(self.left, builder, locals) + left_end_block = builder.block + right_block = builder.append_basic_block() end_block = builder.append_basic_block() - left = ir_to_llvm_expression(self.left, builder, locals) - builder.cbranch(left, end_block, right_block) builder.position_at_end(right_block) right = ir_to_llvm_expression(self.right, builder, locals) + right_end_block = builder.block builder.branch(end_block) builder.position_at_end(end_block) phi = builder.phi(llvm_boolean_type) - phi.add_incoming(llvm.Constant(llvm_boolean_type, 1), left_block) - phi.add_incoming(right, right_block) + phi.add_incoming(llvm.Constant(llvm_boolean_type, 1), left_end_block) + phi.add_incoming(right, right_end_block) return phi diff --git a/tests/test_combinatorically.py b/tests/test_combinatorically.py index 0d29b30..242e95e 100644 --- a/tests/test_combinatorically.py +++ b/tests/test_combinatorically.py @@ -63,6 +63,40 @@ def test_vector_binary(operator, dense1, dense2, format1, format2, format_out): ) +@pytest.mark.parametrize("dense1", [[0, 2, 4, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("dense2", [[-1, 3.5, 0, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("dense3", [[0, 0, 6, -2], [0, 0, 0, 0]]) +@pytest.mark.parametrize("format1", ["s", "d"]) +@pytest.mark.parametrize("format2", ["s", "d"]) +@pytest.mark.parametrize("format3", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["s", "d"]) +def test_vector_add_3(dense1, dense2, dense3, format1, format2, format3, format_out): + assert_same_as_dense( + "out(i) = in1(i) + in2(i) + in3(i)", + format_out, + in1=(dense1, format1), + in2=(dense2, format2), + in3=(dense3, format3), + ) + + +@pytest.mark.parametrize("dense1", [[0, 2, 4, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("dense2", [[-1, 3.5, 0, 0], [0, 0, 0, 0]]) +@pytest.mark.parametrize("dense3", [[0, 0, 6, -2], [0, 0, 0, 0]]) +@pytest.mark.parametrize("format1", ["s", "d"]) +@pytest.mark.parametrize("format2", ["s", "d"]) +@pytest.mark.parametrize("format3", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["s", "d"]) +def test_vector_multiply_3(dense1, dense2, dense3, format1, format2, format3, format_out): + assert_same_as_dense( + "out(i) = in1(i) * in2(i) * in3(i)", + format_out, + in1=(dense1, format1), + in2=(dense2, format2), + in3=(dense3, format3), + ) + + @pytest.mark.parametrize("dense1", [[[0, 2, 4], [0, -1, 0]], [[0, 0, 0], [0, 0, 0]]]) @pytest.mark.parametrize("dense2", [[[-1, 3.5], [0, 0], [4, 0]], [[0, 0], [0, 0], [0, 0]]]) @pytest.mark.parametrize("format1", ["ss", "dd", "sd", "ds", "d1d0"])