[Bugfix] Reorder Passes: Place Vectorize Loop Before StorageFlatten and...

[Bugfix] Reorder Passes: Place Vectorize Loop Before StorageFlatten and FlattenBuffer to Prevent Redundant Allocations (#37)

* installation script fix

* readme typo fix

* doc fix for dequantize gemm

* [Doc] remove CODE_OF_CONDUCT.md and SECURITY.md; update references in CONTRIBUTING.md

* [Doc] add unit tests for AnnotateDeviceRegions transform; remove SUPPORT.md

* update license

* [Enhancement] add tensor supply handling for unsigned integers; improve error message for execution backend assertion

* [Refactor] improve code readability by reformatting function signatures and assertions

* [Refactor] replace torch.manual_seed with tilelang.testing.set_random_seed for consistency in random seed handling

* [Refactor] unify thread binding variable naming across kernel and example files

* [Refactor] remove unused thread binding parameter from matrix multiplication functions

* [Refactor] remove unused thread binding parameter from matrix multiplication functions

* [Refactor] enable main testing function in tilelang kernel gemm test

* bug fix

* lint fix

* [Refactor] reorder vectorize loop

tilelang/engine/lower.py

@@ -144,6 +144,8 @@ def lower(
     mod = tl.transform.LegalizeSafeMemoryAccess()(mod)
     # Inject Simplify to remove the duplicated conditions
     mod = tir.transform.Simplify()(mod)
+    mod = tir.transform.VectorizeLoop()(mod)
+
     # which may be introduced by the LegalizeSafeMemoryAccess
     if target.arch == "sm_90":
         mod = tl.transform.MultiVersionBuffer()(mod)
@@ -161,7 +163,6 @@ def lower(
     mod = tir.transform.FlattenBuffer()(mod)
     mod = tir.transform.NarrowDataType(32)(mod)
     mod = tir.transform.Simplify()(mod)
-    mod = tir.transform.VectorizeLoop()(mod)
     mod = tir.transform.StorageRewrite()(mod)
     mod = tir.transform.UnrollLoop()(mod)
     mod = tir.transform.RenormalizeSplitPattern()(mod)

tilelang/language/parser/parser.py

@@ -463,6 +463,8 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
     elif isinstance(res, str):
         # Ignore docstrings
         pass
+    elif isinstance(res, tvm.tir.stmt.BufferStore):
+        T.buffer_store(res.buffer, res.value, res.indices, res.predicate)
     else:
         self.report_error(node, f"Parsing resulted in unexpected type {type(res)}")
 
@@ -480,13 +482,8 @@ def visit_if(self: Parser, node: doc.If) -> None:
         The doc AST if node.
     """
     with self.var_table.with_frame():
-        condition = self.eval_expr(node.test)
-        if isinstance(condition, bool):
-            if condition:
-                self.visit_body(node.body)
-            elif node.orelse:
-                self.visit_body(node.orelse)
-        else:
+        predicate = self.eval_expr(node.test)
+        if isinstance(predicate, (PrimExpr, tvm.tir.expr.ExprOp)):
             with T.If(self.eval_expr(node.test)):
                 with T.Then():
                     with self.var_table.with_frame():
@@ -495,6 +492,16 @@ def visit_if(self: Parser, node: doc.If) -> None:
                     with T.Else():
                         with self.var_table.with_frame():
                             self.visit_body(node.orelse)
+        elif isinstance(predicate, bool):
+            if predicate:
+                with self.var_table.with_frame():
+                    self.visit_body(node.body)
+            elif node.orelse:
+                with self.var_table.with_frame():
+                    self.visit_body(node.orelse)
+        else:
+            self.report_error(node.test,
+                              f"If condition must be a boolean expression, but got {predicate}")
 
 
 @dispatch.register(token="tir", type_name="Assert")
@@ -529,6 +536,8 @@ def visit_return(self: Parser, node: doc.Return) -> None:
         The doc AST return node.
     """
     value = self.eval_expr(node.value)
+    if value is None:
+        self.report_error(node, "Expression to be returned must be a PrimExpr")
     T.evaluate(tvm.tir.ret(value))