From 507493b436589b8cdda92a6b62a51759218f9762 Mon Sep 17 00:00:00 2001
From: Takeshi Yoneda <t.y.mathetake@gmail.com>
Date: Mon, 10 Jun 2024 12:48:20 -0700
Subject: [PATCH] ssa: optimizes findValue function (#2245)

This removes the unnecessary search and addition of block params.
As a result, the compilation gets faster up to 30% while having no
impacts on the runtime performance.

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
---
 .../engine/wazevo/frontend/frontend_test.go   | 27 +++++------
 internal/engine/wazevo/ssa/builder.go         | 47 ++++++++++++++-----
 2 files changed, 45 insertions(+), 29 deletions(-)

diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index d1040a4a12..b982556d8b 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1058,14 +1058,14 @@ blk1: () <-- (blk0)
 	Call f1:sig1, exec_ctx, module_ctx
 	v5:i64 = Load module_ctx, 0x8
 	v6:i64 = Uload32 module_ctx, 0x10
-	Jump blk3, v2
+	Jump blk3
 
 blk2: () <-- (blk0)
-	Jump blk3, v2
+	Jump blk3
 
-blk3: (v7:i32) <-- (blk1,blk2)
+blk3: () <-- (blk1,blk2)
 	v8:i64 = Iconst_64 0x4
-	v9:i64 = UExtend v7, 32->64
+	v9:i64 = UExtend v2, 32->64
 	v10:i64 = Uload32 module_ctx, 0x10
 	v11:i64 = Iadd v9, v8
 	v12:i32 = Icmp lt_u, v10, v11
@@ -2869,7 +2869,7 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32)
 
 blk1: () <-- (blk0)
 	v11:i32 = Load v9, 0x10
-	Jump blk3, v2
+	Jump blk3
 
 blk2: () <-- (blk0)
 	v12:i32 = Load v9, 0x10
@@ -2880,18 +2880,13 @@ blk2: () <-- (blk0)
 	ExitIfTrue v16, exec_ctx, memory_out_of_bounds
 	v17:i32 = Load v9, 0x30
 	v18:i32 = Load v9, 0x25
-	Jump blk3, v2
+	Jump blk3
 
-blk3: (v19:i32) <-- (blk1,blk2)
-	v20:i64 = Iconst_64 0x19
-	v21:i64 = UExtend v19, 32->64
-	v22:i64 = Uload32 module_ctx, 0x10
-	v23:i64 = Iadd v21, v20
-	v24:i32 = Icmp lt_u, v22, v23
-	ExitIfTrue v24, exec_ctx, memory_out_of_bounds
-	v25:i64 = Load module_ctx, 0x8
-	v26:i64 = Iadd v25, v21
-	v27:i32 = Load v26, 0x15
+blk3: () <-- (blk1,blk2)
+	v20:i64 = Load module_ctx, 0x8
+	v21:i64 = UExtend v2, 32->64
+	v22:i64 = Iadd v20, v21
+	v23:i32 = Load v22, 0x15
 	Jump blk_ret
 `,
 		},
diff --git a/internal/engine/wazevo/ssa/builder.go b/internal/engine/wazevo/ssa/builder.go
index 316e626149..0b700c4b19 100644
--- a/internal/engine/wazevo/ssa/builder.go
+++ b/internal/engine/wazevo/ssa/builder.go
@@ -518,21 +518,42 @@ func (b *builder) findValue(typ Type, variable Variable, blk *basicBlock) Value
 	// If this block has multiple predecessors, we have to gather the definitions,
 	// and treat them as an argument to this block.
 	//
-	// The first thing is to define a new parameter to this block which may or may not be redundant, but
-	// later we eliminate trivial params in an optimization pass. This must be done before finding the
-	// definitions in the predecessors so that we can break the cycle.
-	paramValue := blk.AddParam(b, typ)
-	b.DefineVariable(variable, paramValue, blk)
-
-	// After the new param is added, we have to manipulate the original branching instructions
-	// in predecessors so that they would pass the definition of `variable` as the argument to
-	// the newly added PHI.
+	// But before that, we have to check if the possible definitions are the same Value.
+	tmpValue := b.allocateValue(typ)
+	// Break the cycle by defining the variable with the tmpValue.
+	b.DefineVariable(variable, tmpValue, blk)
+	// Check all the predecessors if they have the same definition.
+	uniqueValue := ValueInvalid
 	for i := range blk.preds {
-		pred := &blk.preds[i]
-		value := b.findValue(typ, variable, pred.blk)
-		pred.branch.addArgumentBranchInst(b, value)
+		predValue := b.findValue(typ, variable, blk.preds[i].blk)
+		if uniqueValue == ValueInvalid {
+			uniqueValue = predValue
+		} else if uniqueValue != predValue {
+			uniqueValue = ValueInvalid
+			break
+		}
+	}
+
+	if uniqueValue != ValueInvalid {
+		// If all the predecessors have the same definition, we can use that value.
+		b.DefineVariable(variable, uniqueValue, blk)
+		b.alias(tmpValue, uniqueValue)
+		return uniqueValue
+	} else {
+		// Otherwise, add the tmpValue to this block as a parameter which may or may not be redundant, but
+		// later we eliminate trivial params in an optimization pass. This must be done before finding the
+		// definitions in the predecessors so that we can break the cycle.
+		blk.addParamOn(tmpValue)
+		// After the new param is added, we have to manipulate the original branching instructions
+		// in predecessors so that they would pass the definition of `variable` as the argument to
+		// the newly added PHI.
+		for i := range blk.preds {
+			pred := &blk.preds[i]
+			value := b.findValue(typ, variable, pred.blk)
+			pred.branch.addArgumentBranchInst(b, value)
+		}
+		return tmpValue
 	}
-	return paramValue
 }
 
 // Seal implements Builder.Seal.