Support allocates when not followed by a sequence statement

With a test to back this case up.

Change-Id: I670809f5ee53b583a15d9b783852dda3089756e9

Author: lhutton1

src/tir/contrib/ethosu/passes.cc

@@ -41,19 +41,21 @@ namespace ethosu {
  * "main" that is being offloaded to the NPU.
  *
  * For example,
- *               allocate {
- *                   extern_call(...)
- *                   allocate {
- *     Before:           extern_call(...)
- *                   }
- *               }
+ * Before:
+ *   allocate {
+ *       extern_call(...)
+ *           allocate {
+ *               extern_call(...)
+ *           }
+ *   }
  *
- *               allocate {
- *                   allocate {
- *                      extern_call(...)
- *     After:           extern_call(...)
- *                   }
- *               }
+ * After:
+ *   allocate {
+ *       allocate {
+ *           extern_call(...)
+ *           extern_call(...)
+ *       }
+ *  }
  */
 class HoistAllocatesMutator : public StmtExprMutator {
  public:
@@ -85,16 +87,17 @@ class HoistAllocatesMutator : public StmtExprMutator {
     allocates_.push_back(GetRef<Allocate>(op));
 
     // Skip the allocate node itself
-    const auto* seq = op->body.as<SeqStmtNode>();
-    ICHECK(seq) << "Expected a sequence statement but got " << op->body->GetTypeKey() << ".";
-
-    // Traverse the allocate body recursively and flatten
-    Array<Stmt> new_stmts;
-    new_stmts.reserve(seq->seq.size());
-    for (const Stmt& old_stmt : seq->seq) {
-      new_stmts.push_back(VisitStmt(old_stmt));
+    if (const auto* seq = op->body.as<SeqStmtNode>()) {
+      // Traverse the allocate body recursively and flatten
+      Array<Stmt> new_stmts;
+      new_stmts.reserve(seq->seq.size());
+      for (const Stmt& old_stmt : seq->seq) {
+        new_stmts.push_back(VisitStmt(old_stmt));
+      }
+      return SeqStmt::Flatten(new_stmts);
+    } else {
+      return VisitStmt(op->body);
     }
-    return SeqStmt::Flatten(new_stmts);
   }
 
   /*! A stack to store allocates as they are visited. */

tests/python/contrib/test_ethosu/test_hoist_allocates.py

@@ -204,6 +204,48 @@ def main(placeholder: T.Buffer[(8192,), "int8"], ethosu_write: T.Buffer[(2048,),
     CheckAllocates(allocate_info)(mod)
 
 
+def test_allocate_without_seq_stmt():
+    """
+    Tests the case when an allocate statement does not have a sequence statement as its body.
+    """
+    # fmt: off
+    @tvm.script.ir_module
+    class Module:
+        @T.prim_func
+        def main(placeholder: T.Buffer[(8192,), "int8"], ethosu_write: T.Buffer[(2048,), "int8"], buffer_encoded: T.Buffer[(128,), "uint8"], buffer_encoded_1: T.Buffer[(32,), "uint8"], buffer_encoded_2: T.Buffer[(112,), "uint8"], buffer_encoded_3: T.Buffer[(32,), "uint8"], buffer_encoded_4: T.Buffer[(112,), "uint8"], buffer_encoded_5: T.Buffer[(32,), "uint8"], buffer_encoded_6: T.Buffer[(112,), "uint8"], buffer_encoded_7: T.Buffer[(32,), "uint8"]) -> None:
+            # function attr dict
+            T.func_attr({"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True})
+            T.preflattened_buffer(placeholder, [1, 16, 16, 32], dtype="int8", data=placeholder.data)
+            T.preflattened_buffer(ethosu_write, [1, 16, 16, 8], dtype="int8", data=ethosu_write.data)
+            # body
+            placeholder_global = T.allocate([128], "uint8", "global")
+            placeholder_global_1 = T.allocate([112], "uint8", "global")
+            placeholder_global_2 = T.allocate([112], "uint8", "global")
+            placeholder_d_global = T.allocate([32], "uint8", "global")
+            placeholder_d_global_1 = T.allocate([32], "uint8", "global")
+            placeholder_d_global_2 = T.allocate([32], "uint8", "global")
+            placeholder_global_3 = T.allocate([112], "uint8", "global")
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded[0], 128, placeholder_global[0], dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded_1[0], 32, placeholder_d_global[0], dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_conv2d", "int8", 16, 16, 32, 16, 0, 16, placeholder[0], 0, 0, 0, T.float32(0.5), 10, "NHWC", 512, 32, 1, "int8", 16, 16, 2, 16, 0, 16, ethosu_write[0], 0, 0, 0, T.float32(0.25), 14, "NHWC", 128, 8, 1, 1, 1, 1, 1, 1, 1, placeholder_global[0], 128, 12, placeholder_d_global[0], 32, 0, 0, 0, 0, "NONE", 0, 0, "TFL", "NONE", 0, 0, 0, dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded_2[0], 112, placeholder_global_1[0], dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded_3[0], 32, placeholder_d_global_1[0], dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_conv2d", "int8", 16, 16, 32, 16, 0, 16, placeholder[0], 0, 0, 0, T.float32(0.5), 10, "NHWC", 512, 32, 1, "int8", 16, 16, 2, 16, 0, 16, ethosu_write[2], 0, 0, 0, T.float32(0.25), 14, "NHWC", 128, 8, 1, 1, 1, 1, 1, 1, 1, placeholder_global_1[0], 112, 12, placeholder_d_global_1[0], 32, 0, 0, 0, 0, "NONE", 0, 0, "TFL", "NONE", 0, 0, 0, dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded_4[0], 112, placeholder_global_2[0], dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded_5[0], 32, placeholder_d_global_2[0], dtype="handle"))
+            placeholder_d_global_3 = T.allocate([32], "uint8", "global")
+            T.evaluate(T.call_extern("ethosu_conv2d", "int8", 16, 16, 32, 16, 0, 16, placeholder[0], 0, 0, 0, T.float32(0.5), 10, "NHWC", 512, 32, 1, "int8", 16, 16, 2, 16, 0, 16, ethosu_write[4], 0, 0, 0, T.float32(0.25), 14, "NHWC", 128, 8, 1, 1, 1, 1, 1, 1, 1, placeholder_global_2[0], 112, 12, placeholder_d_global_2[0], 32, 0, 0, 0, 0, "NONE", 0, 0, "TFL", "NONE", 0, 0, 0, dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded_6[0], 112, placeholder_global_3[0], dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_copy", buffer_encoded_7[0], 32, placeholder_d_global_3[0], dtype="handle"))
+            T.evaluate(T.call_extern("ethosu_conv2d", "int8", 16, 16, 32, 16, 0, 16, placeholder[0], 0, 0, 0, T.float32(0.5), 10, "NHWC", 512, 32, 1, "int8", 16, 16, 2, 16, 0, 16, ethosu_write[6], 0, 0, 0, T.float32(0.25), 14, "NHWC", 128, 8, 1, 1, 1, 1, 1, 1, 1, placeholder_global_3[0], 112, 12, placeholder_d_global_3[0], 32, 0, 0, 0, 0, "NONE", 0, 0, "TFL", "NONE", 0, 0, 0, dtype="handle"))
+    # fmt: on
+
+    mod = Module
+    allocate_info = ExtractAllocateInfo()(mod)
+    mod = HoistAllocates()(mod)
+    CheckAllocates(allocate_info)(mod)
+
+
 def test_multiple_prim_funcs():
     @tvm.script.ir_module
     class Module: