[Fixbug][Pass] Fix a bug in the inline_let_stmt pass (#240)

.
hidet-org · May 23, 2023 · e198cd1 · e198cd1
1 parent 7384232
commit e198cd1
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/hidet/transforms/inline_let_stmt.py b/python/hidet/transforms/inline_let_stmt.py
@@ -79,11 +79,12 @@ def visit_LetStmt(self, stmt: LetStmt):
         bind_vars = []
         bind_values = []
         for bind_var, bind_value in zip(stmt.bind_vars, stmt.bind_values):
+            updated_var = self(bind_var)
             updated_value = self(bind_value)
-            if self.should_inline(bind_var, updated_value):
+            if self.should_inline(updated_var, updated_value):
                 self.memo[bind_var] = updated_value
             else:
-                bind_vars.append(bind_var)
+                bind_vars.append(updated_var)
                 bind_values.append(updated_value)
         body = self(stmt.body)
         if same_list(bind_vars, stmt.bind_vars) and same_list(bind_values, stmt.bind_values) and body is stmt.body: