[Enhancement] Add role assignment for AllocateNode in warp specialization (#657)

- Implemented a new role assignment for `AllocateNode` in `warp_specialized_rewriter.cc`, setting the role to `kConsumer` to ensure proper handling of memory allocation scenarios.
- This can avoid bug when using T.reduce(clear=False)

Author: chengyupku

src/transform/warp_specialized_rewriter.cc

@@ -170,6 +170,12 @@ class WarpSpecializedRoleMarker : public StmtVisitor {
     SetRole(op, GetRole(op->block));
   }
 
+  void VisitStmt_(const AllocateNode *op) final {
+    StmtVisitor::VisitStmt_(op);
+    Role role = Role::kConsumer;
+    SetRole(op, role);
+  }
+
   template <class NodeType> void HandleBodyStmt(const NodeType *op) {
     StmtVisitor::VisitStmt_(op);
     SetRole(op, GetRole(op->body));