diff --git a/mace/proto/mace.proto b/mace/proto/mace.proto
index 8e6808219ce27cbe4f122a1e4b909c71f7bab7a3..7a0ad7211737a8efd41eafa71b9a2c07ed2815da 100644
--- a/mace/proto/mace.proto
+++ b/mace/proto/mace.proto
@@ -67,12 +67,20 @@ message NodeInput {
   optional int32 output_port = 2;
 }
 
+message OutputShape {
+  repeated int64 dims = 1;
+}
+
 message OperatorDef {
   repeated string input = 1;
   repeated string output = 2;
   optional string name = 3;
   optional string type = 4;
   repeated Argument arg = 5;
+  optional OutputShape output_shape = 6;
+
+  // Memory optimization: only support one single output op
+  optional int32 mem_id = 10 [default = -1];
 
   // for hexagon mace-nnlib
   optional uint32 node_id = 100;
@@ -82,6 +90,16 @@ message OperatorDef {
   repeated int32 out_max_byte_size = 104; // only support 32-bit len
 }
 
+// for memory optimization
+message MemoryBlock {
+  optional int32 mem_id = 1;
+  optional uint32 x = 2;
+  optional uint32 y = 3;
+}
+message MemoryArena {
+  repeated MemoryBlock mem_block = 1;
+}
+
 // for hexagon mace-nnlib
 message InputInfo {
   optional string name = 1;