dnnobf.patch

diff --git a/include/tvm/ir/attrs.h b/include/tvm/ir/attrs.h
index 18d0f025c..5ed85ec71 100644
--- a/include/tvm/ir/attrs.h
+++ b/include/tvm/ir/attrs.h
@@ -298,6 +298,11 @@ class DictAttrs : public Attrs {
     return GetAttr<Integer>(attr_key, 0).value_or(0).IntValue() != 0;
   }
 
+  // zz: for obfuscation
+  bool HasIntAttr(const std::string& attr_key) const {
+    return GetAttr<Integer>(attr_key, 0).defined();
+  }
+
   TVM_DEFINE_OBJECT_REF_METHODS(DictAttrs, Attrs, DictAttrsNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(DictAttrsNode);
 };
diff --git a/include/tvm/relay/expr.h b/include/tvm/relay/expr.h
index 854050464..f30398fca 100644
--- a/include/tvm/relay/expr.h
+++ b/include/tvm/relay/expr.h
@@ -300,6 +300,9 @@ class CallNode : public ExprNode {
   /*! \brief The additional attributes */
   Attrs attrs;
 
+  // zz: add attrs for obfuscation
+  DictAttrs obf_attrs;
+
   /*!
    * \brief The type arguments passed to polymorphic(template) function.
    *
@@ -324,6 +327,7 @@ class CallNode : public ExprNode {
     v->Visit("op", &op);
     v->Visit("args", &args);
     v->Visit("attrs", &attrs);
+    v->Visit("obf_attrs", &obf_attrs); // zz: add obf_attrs
     v->Visit("type_args", &type_args);
     v->Visit("virtual_device_", &virtual_device_);
     v->Visit("span", &span);
@@ -342,6 +346,7 @@ class CallNode : public ExprNode {
     hash_reduce(op);
     hash_reduce(args);
     hash_reduce(attrs);
+    hash_reduce(obf_attrs); // zz: add obf_attrs
     if (!IsPrimitiveOp(op)) {
       hash_reduce(type_args);
     }
@@ -370,7 +375,8 @@ class Call : public Expr {
    * \param span The source span of the expression.
    */
   TVM_DLL Call(Expr op, Array<Expr> args, Attrs attrs = Attrs(),
-               Array<Type> type_args = Array<Type>(), Span span = Span());
+              //  Array<Type> type_args = Array<Type>(), Span span = Span());
+                Array<Type> type_args = Array<Type>(), Span span = Span(), DictAttrs obf_attrs = DictAttrs()); // zz: add obf_attrs
 
   TVM_DEFINE_OBJECT_REF_METHODS(Call, RelayExpr, CallNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
@@ -386,7 +392,20 @@ Call WithFields(Call call, Optional<Expr> opt_op = Optional<Expr>(),
                 Optional<Attrs> opt_attrs = Optional<Attrs>(),
                 Optional<Array<Type>> opt_type_args = Optional<Array<Type>>(),
                 Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
-                Optional<Span> opt_span = Optional<Span>());
+                // Optional<Span> opt_span = Optional<Span>());
+                Optional<Span> opt_span = Optional<Span>(), Optional<DictAttrs> opt_obf_attrs = Optional<DictAttrs>()); // zz: add obf_attrs
+
+// zz: for obfuscation
+Call WithObfAttr(Call input, const std::string& attr_key, ObjectRef attr_value);
+// zz: ref include/tvm/ir/module.h GetAttr
+template <typename TObjectRef>
+Optional<TObjectRef> GetObfAttr(
+    Call call,
+    const std::string& attr_key,
+    Optional<TObjectRef> default_value = Optional<TObjectRef>(nullptr)) {
+  return call->obf_attrs.GetAttr<ObjectRef>(attr_key, default_value);
+}
+
 
 /*!
  * \brief Let binding that binds a local var and optionally a type annotation.
diff --git a/include/tvm/relay/function.h b/include/tvm/relay/function.h
index 874d4f233..ef240d88c 100644
--- a/include/tvm/relay/function.h
+++ b/include/tvm/relay/function.h
@@ -172,6 +172,11 @@ constexpr const char* kPrimitive = "Primitive";
  */
 constexpr const char* kExtern = "Extern";
 
+/*!
+  added by zz: indicate function to be obfuscated
+ */
+constexpr const char* kObf = "Obf";
+
 /*!
  * \brief Indicates the name of the external codegen 'compiler' that should be used to lower
  * or compile the function other than TVM's default lowering pipeline. The name may correspond
diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index 35f0852f3..bba2ccaf1 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -146,7 +146,14 @@ TVM_DLL Pass SplitArgs(uint64_t max_function_args);
  *
  * \return The pass.
  */
-TVM_DLL Pass FuseOps(int fuse_opt_level = -1);
+TVM_DLL Pass FuseOps(int fuse_opt_level = -1, bool enable_flex_fuse = false, bool enable_fakeop_insert = false); // zz: add obf switches
+// added by zz: for dnnobuse impl
+// TVM_DLL Pass FuseOps(int fuse_opt_level = -1, bool skip_fuse = false);
+// TVM_DLL Pass FuseOps_faked(int fuse_opt_level = -1);
+TVM_DLL Pass FuseOpsObf(int fuse_opt_level = -1, bool enable_flex_fuse = true, bool enable_fakeop_insert = true); // zz: obf switches
+
+TVM_DLL Pass AnnotateObfAttrs(); // zz: neuron range annotation for obfuscation
+TVM_DLL Pass InterestingExprExaminer(); // zz: for neuron range profiling
 
 /*!
  * \brief The inverse operation of FuseOps. It transforms a fused program returned by
@@ -591,6 +598,9 @@ TVM_DLL Pass FlattenAtrousConv();
  */
 TVM_DLL Pass AnnotateUsedMemory();
 
+// zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+TVM_DLL Pass AnnotateUsedMemoryWithMainFuncName(String main_func_name);
+
 /*!
  * \brief Captures the post-dfs index and dominator post-dfs index of (most) expression nodes in
  * their span, in the form "index:<post-dfs index>:<dominator post-dfs index>". This is useful for
@@ -618,6 +628,9 @@ TVM_DLL Pass AnnotateMemoryScope();
  */
 TVM_DLL Pass RemoveStandaloneReshapes();
 
+// zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+TVM_DLL Pass RemoveStandaloneReshapesWithMainFuncName(String main_func_name);
+
 }  // namespace transform
 
 /*!
diff --git a/include/tvm/runtime/metadata.h b/include/tvm/runtime/metadata.h
index f921f3e39..a44a80a1e 100644
--- a/include/tvm/runtime/metadata.h
+++ b/include/tvm/runtime/metadata.h
@@ -69,6 +69,9 @@ class MetadataNode : public MetadataBaseNode {
   inline int64_t num_constant_pools() const { return data_->num_constant_pools; }
   TVM_DECLARE_FINAL_OBJECT_INFO(MetadataNode, MetadataBaseNode);
 
+  // zz: for debug purpose (see src/target/metadata_module.cc)
+  const struct ::TVMMetadata* get_metadata() { return data_; }
+
  private:
   const struct ::TVMMetadata* data_;
 };
diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h
index 8400344bf..4a06d1514 100644
--- a/include/tvm/runtime/ndarray.h
+++ b/include/tvm/runtime/ndarray.h
@@ -125,6 +125,8 @@ class NDArray : public ObjectRef {
    * \param stream The output data stream
    */
   inline void Save(dmlc::Stream* stream) const;
+  // zz: for save metadata, i.e., params in aot build (check src/target/metadata_module.cc, ref: Save())
+  inline size_t SaveData(dmlc::Stream* stream) const;
   /*!
    * \brief Create a NDArray that shares the data memory with the current one.
    * \param shape The shape of the new array.
@@ -440,10 +442,13 @@ inline Object* TVMArrayHandleToObjectHandle(TVMArrayHandle handle) {
 /*! \brief Magic number for NDArray file */
 constexpr uint64_t kTVMNDArrayMagic = 0xDD5E40F096B4A13F;
 
+// zz: add some debug info in this function for tensor data inspection (check src/target/metadata.h InMemoryMetadataNode)
 inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor) {
   uint64_t header = kTVMNDArrayMagic, reserved = 0;
   strm->Write(header);
+  std::cout << "  (zz) SaveDLTensor Write " << sizeof(header) << " bytes header" << std::endl;
   strm->Write(reserved);
+  std::cout << "  (zz) SaveDLTensor Write " << sizeof(reserved) << " bytes reserved" << std::endl;
   // Always save data as CPU context
   //
   // Parameters that get serialized should be in CPU by default.
@@ -457,10 +462,14 @@ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor) {
   cpu_dev.device_type = kDLCPU;
   cpu_dev.device_id = 0;
   strm->Write(cpu_dev);
+  std::cout << "  (zz) SaveDLTensor Write " << sizeof(cpu_dev) << " bytes cpu_dev" << std::endl;
   strm->Write(tensor->ndim);
+  std::cout << "  (zz) SaveDLTensor Write " << sizeof(tensor->ndim) << " bytes tensor->ndim" << std::endl;
   strm->Write(tensor->dtype);
+  std::cout << "  (zz) SaveDLTensor Write " << sizeof(tensor->dtype) << " bytes tensor->dtype" << std::endl;
   int ndim = tensor->ndim;
   strm->WriteArray(tensor->shape, ndim);
+  std::cout << "  (zz) SaveDLTensor Write " << sizeof(*tensor->shape) * ndim << " bytes tensor->shape" << std::endl;
   int type_bytes = (tensor->dtype.bits + 7) / 8;
   int64_t num_elems = 1;
   for (int i = 0; i < ndim; ++i) {
@@ -468,11 +477,13 @@ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor) {
   }
   int64_t data_byte_size = type_bytes * num_elems;
   strm->Write(data_byte_size);
+  std::cout << "  (zz) SaveDLTensor Write " << sizeof(data_byte_size) << " bytes data_byte_size" << std::endl;
 
   if (DMLC_IO_NO_ENDIAN_SWAP && tensor->device.device_type == kDLCPU &&
       tensor->strides == nullptr && tensor->byte_offset == 0) {
     // quick path
     strm->Write(tensor->data, data_byte_size);
+    std::cout << "  (zz) SaveDLTensor Write " << data_byte_size << " bytes tensor->data" << std::endl;
   } else {
     std::vector<uint8_t> bytes(data_byte_size);
     ICHECK_EQ(
@@ -483,12 +494,67 @@ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor) {
       dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
     }
     strm->Write(dmlc::BeginPtr(bytes), data_byte_size);
+    std::cout << "  (zz) SaveDLTensor Write " << data_byte_size << " bytes dmlc::BeginPtr(bytes) (not in example)" << std::endl;
   }
   return true;
 }
 
+// zz: for save metadata, i.e., params in aot build (check src/target/metadata_module.cc, ref: SaveDLTensor())
+// in this function we only save tensor elements seuqentially (SaveDLTensor also save metadata, e.g., ndim, shape, dtype, etc.)
+inline size_t SaveDLTensorData(dmlc::Stream* strm, const DLTensor* tensor) {
+  // uint64_t header = kTVMNDArrayMagic, reserved = 0;
+  // strm->Write(header);
+  // std::cout << "   SaveDLTensor Write " << sizeof(header) << " bytes header" << std::endl;
+  // strm->Write(reserved);
+  // std::cout << "   SaveDLTensor Write " << sizeof(reserved) << " bytes reserved" << std::endl;
+
+  // Device cpu_dev;
+  // cpu_dev.device_type = kDLCPU;
+  // cpu_dev.device_id = 0;
+  // strm->Write(cpu_dev);
+  // std::cout << "   SaveDLTensor Write " << sizeof(cpu_dev) << " bytes cpu_dev" << std::endl;
+  // strm->Write(tensor->ndim);
+  // std::cout << "   SaveDLTensor Write " << sizeof(tensor->ndim) << " bytes tensor->ndim" << std::endl;
+  // strm->Write(tensor->dtype);
+  // std::cout << "   SaveDLTensor Write " << sizeof(tensor->dtype) << " bytes tensor->dtype" << std::endl;
+  int ndim = tensor->ndim;
+  // strm->WriteArray(tensor->shape, ndim);
+  // std::cout << "   SaveDLTensor Write " << sizeof(*tensor->shape) * ndim << " bytes tensor->shape" << std::endl;
+  int type_bytes = (tensor->dtype.bits + 7) / 8;
+  int64_t num_elems = 1;
+  for (int i = 0; i < ndim; ++i) {
+    num_elems *= tensor->shape[i];
+  }
+  int64_t data_byte_size = type_bytes * num_elems;
+  // strm->Write(data_byte_size);
+  // std::cout << "   SaveDLTensor Write " << sizeof(data_byte_size) << " bytes data_byte_size" << std::endl;
+
+  if (DMLC_IO_NO_ENDIAN_SWAP && tensor->device.device_type == kDLCPU &&
+      tensor->strides == nullptr && tensor->byte_offset == 0) {
+    // quick path
+    strm->Write(tensor->data, data_byte_size);
+    // std::cout << "   SaveDLTensor Write " << data_byte_size << " bytes tensor->data" << std::endl;
+  } else {
+    std::vector<uint8_t> bytes(data_byte_size);
+    ICHECK_EQ(
+        TVMArrayCopyToBytes(const_cast<DLTensor*>(tensor), dmlc::BeginPtr(bytes), data_byte_size),
+        0)
+        << TVMGetLastError();
+    if (!DMLC_IO_NO_ENDIAN_SWAP) {
+      dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
+    }
+    strm->Write(dmlc::BeginPtr(bytes), data_byte_size);
+    // std::cout << "   SaveDLTensor Write " << data_byte_size << " bytes dmlc::BeginPtr(bytes) (not in lenet)" << std::endl;
+  }
+  // return true;
+  return data_byte_size;
+}
+
 inline void NDArray::Save(dmlc::Stream* strm) const { SaveDLTensor(strm, operator->()); }
 
+// zz: for save metadata, i.e., params in aot build (check src/target/metadata_module.cc, ref: Save())
+inline size_t NDArray::SaveData(dmlc::Stream* strm) const { return SaveDLTensorData(strm, operator->()); }
+
 inline bool NDArray::Load(dmlc::Stream* strm) {
   uint64_t header, reserved;
   ICHECK(strm->Read(&header)) << "Invalid DLTensor file format";
diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h
index a1697d807..2b2dac340 100644
--- a/include/tvm/tir/transform.h
+++ b/include/tvm/tir/transform.h
@@ -627,6 +627,9 @@ TVM_DLL Pass ConvertForLoopsToSerial();
  */
 TVM_DLL Pass UnifiedStaticMemoryPlanner();
 
+// zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+TVM_DLL Pass UnifiedStaticMemoryPlannerWithMainFuncName(String main_func_name);
+
 /*!
  * \brief This pass transforms annotated loops into pipelined ones where producers and consumers
  * are overlapped with the information provided in loop annotations, which enables optimization
diff --git a/include/tvm/tir/usmp/transform.h b/include/tvm/tir/usmp/transform.h
index ccb684463..eb701e492 100644
--- a/include/tvm/tir/usmp/transform.h
+++ b/include/tvm/tir/usmp/transform.h
@@ -46,6 +46,10 @@ using Pass = tvm::transform::Pass;
 TVM_DLL Pass ConvertPoolAllocationsToOffsets(const Map<tir::Stmt, PoolAllocation>& pool_allocations,
                                              Bool emit_tvmscript_printable = Bool(false));
 
+// zz: for obfuscated function, used in PlanMemoryWithMainFuncName
+TVM_DLL Pass ConvertPoolAllocationsToOffsetsWithMainFuncName(const Map<tir::Stmt, PoolAllocation>& pool_allocations,
+                                             String main_func_name, Bool emit_tvmscript_printable = Bool(false));
+
 /*!
  * \brief Assign PoolInfo objects to tir.allocate nodes depending on the PrimFunc's target
  *
@@ -56,6 +60,9 @@ TVM_DLL Pass ConvertPoolAllocationsToOffsets(const Map<tir::Stmt, PoolAllocation
  */
 TVM_DLL Pass AssignPoolInfo();
 
+// zz: for obfuscated function, used in PlanMemoryWithMainFuncName
+TVM_DLL Pass AssignPoolInfoWithMainFuncName(String main_func_name);
+
 /*!
  * \brief This pass creates Allocate nodes for I/O tensors
  *
diff --git a/include/tvm/tir/usmp/utils.h b/include/tvm/tir/usmp/utils.h
index a67350a2b..e0a7c2cb9 100644
--- a/include/tvm/tir/usmp/utils.h
+++ b/include/tvm/tir/usmp/utils.h
@@ -250,6 +250,9 @@ Array<BufferInfo> ConvertToArrayOfBufferInfo(const Map<BufferInfo, Stmt>& buffer
  */
 Integer CalculateModuleWorkspaceSize(const IRModule& mod);
 
+// zz: for obfuscated function, used in PlanMemoryWithMainFuncName
+Integer CalculateModuleWorkspaceSizeWithMainFuncName(const IRModule& mod, String main_func_name);
+
 /*!
  * \brief The allocate node attribute to indicate candidate memory pools.
  * This needs to be kept in sync with CANDIDATE_MEMORY_POOL_ATTR in
diff --git a/python/tvm/relay/expr.py b/python/tvm/relay/expr.py
index 5239eaa88..cb5ad96ff 100644
--- a/python/tvm/relay/expr.py
+++ b/python/tvm/relay/expr.py
@@ -344,6 +344,12 @@ def CallWithFields(
     """
     return _ffi_api.CallWithFields(call, op, args, attrs, type_args, virtual_device, span)
 
+# zz: for obfuscation
+def CallWithObfAttr(call, attr_key, attr_value):
+    return _ffi_api.CallWithObfAttr(call, attr_key, tvm.runtime.convert(attr_value))
+
+def CallGetObfAttr(call, attr_key):
+    return _ffi_api.CallGetObfAttr(call, attr_key)
 
 @tvm._ffi.register_object("relay.Let")
 class Let(ExprWithOp):
diff --git a/src/relay/analysis/graph_partitioner.cc b/src/relay/analysis/graph_partitioner.cc
index d233d43ad..6f7c32e35 100644
--- a/src/relay/analysis/graph_partitioner.cc
+++ b/src/relay/analysis/graph_partitioner.cc
@@ -24,6 +24,27 @@
 namespace tvm {
 namespace relay {
 
+// zz: helper function
+bool isComputeOp(const CallNode* call_node) {
+  if (auto call_op = call_node->op.as<OpNode>()) {
+    // check if call_op is conv2d, dense, max_pool2d
+    if (call_op->name == "nn.conv2d" || call_op->name == "nn.dense" || call_op->name == "nn.max_pool2d" || call_op->name == "nn.batch_matmul") {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool isSmallComputeOp(const CallNode* call_node) {
+  if (auto call_op = call_node->op.as<OpNode>()) {
+    // check if call_op is conv2d, dense, max_pool2d
+    if (call_op->name == "add" || call_op->name == "nn.relu") {
+      return true;
+    }
+  }
+  return false;
+}
+
 DominatorTree DominatorTree::PostDom(support::Arena* arena, const IndexedForwardGraph& graph) {
   DominatorTree tree;
   tree.nodes.resize(graph.post_dfs_order.size(), nullptr);
@@ -104,10 +125,200 @@ std::vector<GraphPartitioner::Group*> GraphPartitioner::Partition(
   if (opt_level_ == 0) return std::move(groups_);
   // get post dominator tree
   auto post_dom_tree = DominatorTree::PostDom(arena_, graph);
+  // zz: original fusion
   // run fusion algorithm.
   for (int phase = 0; phase < 3; ++phase) {
     this->RunFuse(graph, post_dom_tree, phase);
   }
+
+  // zz: test flexible fusion
+  // std::cout << "[ZZDEBUG] GraphPartitioner::Partition before RunFuseObfFlex" << std::endl;
+  // this->RunFuse(graph, post_dom_tree, 0);
+  // this->RunFuseObfFlex(graph, post_dom_tree); // zz: test flexible fusion for obfuscation
+  // std::cout << "[ZZDEBUG] GraphPartitioner::Partition after RunFuseObfFlex" << std::endl;
+
+  // zz: print group summary
+  // std::cout << "==========================================================================================================================" << std::endl;
+  // std::cout << std::endl << "[ZZDEBUG] RunFuse Summary:" << std::endl;
+  ICHECK(groups_.size() == graph.post_dfs_order.size());
+  ICHECK(groups_.size() == post_dom_tree.nodes.size());
+  for (size_t i = 0; i < groups_.size(); i++) {
+    ICHECK(groups_[i]->root_ref == graph.post_dfs_order[i]->ref);
+    ICHECK(post_dom_tree.nodes[i]->gnode == graph.post_dfs_order[i]);
+    // std::cout << "------------------ nid " << i << std::endl;
+    // std::cout << "node ir: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref))) << std::endl;
+    // std::cout << "node type key: " << graph.post_dfs_order[i]->ref->GetTypeKey() << std::endl;
+    // std::cout << "node tensor type: " << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref))->checked_type()) << std::endl;
+    // std::cout << "#nodes: " << groups_[i]->num_nodes << std::endl;
+    // std::cout << "#args: " << groups_[i]->args_num << std::endl;
+    // std::cout << "#compute nodes: " << groups_[i]->num_compute_nodes << std::endl;
+    auto root_group = groups_[i]->FindRoot();
+    // find root_group index in groups_
+    size_t root_group_index = 0xffffffff;
+    for (size_t j = 0; j < groups_.size(); j++) {
+      if (groups_[j] == root_group) {
+        root_group_index = j;
+        break;
+      }
+    }
+    // std::cout << "root nid: " << root_group_index << std::endl;
+  }
+
+  return std::move(groups_);
+}
+
+// zz: for random fuse testing, split Partition() into PartitionWithoutMov() and GetGroups(), ref Partition()
+void GraphPartitioner::PartitionWithoutMov(
+    const IndexedForwardGraph& graph) {
+  // ----------------------------------------------------------------------------------
+  // zz: for obf flex fuse
+  // ----------------------------------------------------------------------------------
+  this->InitGroups(graph);
+  if (opt_level_ == 0) return;
+  // get post dominator tree
+  auto post_dom_tree = DominatorTree::PostDom(arena_, graph);
+  // run fusion algorithm.
+  // for (size_t i = 0; i < groups_.size(); i++) {
+  //   std::cout << "    groups_[" << i << "]: group_id = " << std::hex << (uintptr_t)groups_[i] << ", root_id = " << (uintptr_t)(groups_[i]->FindRoot()) << ", num_nodes = " << std::dec << groups_[i]->num_nodes << std::endl;
+  //   std::cout << std::endl << "        group node: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref))) << std::endl << "        group node end." << std::endl;
+  // }
+  // std::cout << std::endl << "[ZZDEBUG] GraphPartitioner::Partition: max_fuse_depth_ = " << std::dec << max_fuse_depth_ << ", max_function_args_ = " << max_function_args_ << std::endl;
+
+  // zz: original fuse
+  // for (int phase = 0; phase < 3; ++phase) {
+  //   this->RunFuse(graph, post_dom_tree, phase);
+  // }
+
+  // zz: flexible fuse for obfuscation
+  // std::cout << "[ZZDEBUG] GraphPartitioner::Partition before RunFuseObfFlex" << std::endl;
+  this->RunFuse(graph, post_dom_tree, 0); // resue phase 0 for 0-dim scalar op attr fusion, e.g., nn.pad(%Input3, 0f, ...). (Remove this for transformer to avoid bug) (TODO: make this configurable)
+  this->RunFuseObfFlex(graph, post_dom_tree); // flexible fuse for obfuscation
+  // std::cout << "[ZZDEBUG] GraphPartitioner::Partition after RunFuseObfFlex" << std::endl;
+
+  // zz: print group summary
+  // std::cout << "==========================================================================================================================" << std::endl;
+  // std::cout << std::endl << "[ZZDEBUG] RunFuse Summary:" << std::endl;
+  ICHECK(groups_.size() == graph.post_dfs_order.size());
+  ICHECK(groups_.size() == post_dom_tree.nodes.size());
+  for (size_t i = 0; i < groups_.size(); i++) {
+    ICHECK(groups_[i]->root_ref == graph.post_dfs_order[i]->ref);
+    ICHECK(post_dom_tree.nodes[i]->gnode == graph.post_dfs_order[i]);
+    // std::cout << "------------------ nid " << i << std::endl;
+    // std::cout << "node ir: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref))) << std::endl;
+    Expr node_expr = GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref));
+    if (node_expr->IsInstance<CallNode>()) {
+      auto call_node = node_expr.as<CallNode>();
+      // std::cout << "node op: " << call_node->op << std::endl;
+    }
+    // std::cout << "node type key: " << graph.post_dfs_order[i]->ref->GetTypeKey() << std::endl;
+    // std::cout << "node tensor type: " << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref))->checked_type()) << std::endl;
+    // std::cout << "#nodes: " << groups_[i]->num_nodes << std::endl;
+    // std::cout << "#args: " << groups_[i]->args_num << std::endl;
+    // std::cout << "#compute nodes: " << groups_[i]->num_compute_nodes << std::endl;
+    // find dominator tree node idx
+    auto dom_tree_parent = post_dom_tree.nodes[i]->parent;
+    if (dom_tree_parent != nullptr) {
+      // find index of dom_tree_node
+      auto parent_node = dom_tree_parent->gnode;
+      size_t dom_tree_node_index = 0xffffffff;
+      for (size_t j = 0; j < groups_.size(); j++) {
+        if (groups_[j]->root_ref == parent_node->ref) {
+          dom_tree_node_index = j;
+          break;
+        }
+      }
+      // std::cout << "dom tree node nid: " << dom_tree_node_index << std::endl;
+    } else {
+      // std::cout << "dom tree node: nullptr" << std::endl;
+    }
+    auto root_group = groups_[i]->FindRoot();
+    // find root_group index in groups_
+    size_t root_group_index = 0xffffffff;
+    for (size_t j = 0; j < groups_.size(); j++) {
+      if (groups_[j] == root_group) {
+        root_group_index = j;
+        break;
+      }
+    }
+    // std::cout << "root nid: " << root_group_index << std::endl;
+  }
+
+  // zz: select fake ops from partitioned groups
+  this->SelectFops(graph, post_dom_tree, 0);
+  this->SelectFops(graph, post_dom_tree, 1);
+
+  // return std::move(groups_);
+  // ----------------------------------------------------------------------------------
+  // zz: for obf flex fuse end
+  // ----------------------------------------------------------------------------------
+
+  // ----------------------------------------------------------------------------------
+  // zz: for manual fusion
+  // ----------------------------------------------------------------------------------
+  // this->InitGroups(graph);
+  // if (opt_level_ == 0) return;
+  // // get post dominator tree
+  // auto post_dom_tree = DominatorTree::PostDom(arena_, graph);
+  // // zz: original fusion
+  // // run fusion algorithm.
+  // for (int phase = 0; phase < 3; ++phase) {
+  //   this->RunFuse(graph, post_dom_tree, phase);
+  // }
+
+  // // zz: print group summary
+  // std::cout << "==========================================================================================================================" << std::endl;
+  // std::cout << std::endl << "[ZZDEBUG] RunFuse Summary:" << std::endl;
+  // ICHECK(groups_.size() == graph.post_dfs_order.size());
+  // ICHECK(groups_.size() == post_dom_tree.nodes.size());
+  // for (size_t i = 0; i < groups_.size(); i++) {
+  //   ICHECK(groups_[i]->root_ref == graph.post_dfs_order[i]->ref);
+  //   ICHECK(post_dom_tree.nodes[i]->gnode == graph.post_dfs_order[i]);
+  //   std::cout << "------------------ nid " << i << std::endl;
+  //   // std::cout << "node ir: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref))) << std::endl;
+  //   Expr node_expr = GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref));
+  //   if (node_expr->IsInstance<CallNode>()) {
+  //     auto call_node = node_expr.as<CallNode>();
+  //     std::cout << "node op: " << call_node->op << std::endl;
+  //   }
+  //   std::cout << "node type key: " << graph.post_dfs_order[i]->ref->GetTypeKey() << std::endl;
+  //   std::cout << "node tensor type: " << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph.post_dfs_order[i]->ref))->checked_type()) << std::endl;
+  //   std::cout << "#nodes: " << groups_[i]->num_nodes << std::endl;
+  //   std::cout << "#args: " << groups_[i]->args_num << std::endl;
+  //   std::cout << "#compute nodes: " << groups_[i]->num_compute_nodes << std::endl;
+  //   // find dominator tree node idx
+  //   auto dom_tree_parent = post_dom_tree.nodes[i]->parent;
+  //   if (dom_tree_parent != nullptr) {
+  //     // find index of dom_tree_node
+  //     auto parent_node = dom_tree_parent->gnode;
+  //     size_t dom_tree_node_index = 0xffffffff;
+  //     for (size_t j = 0; j < groups_.size(); j++) {
+  //       if (groups_[j]->root_ref == parent_node->ref) {
+  //         dom_tree_node_index = j;
+  //         break;
+  //       }
+  //     }
+  //     std::cout << "dom tree node nid: " << dom_tree_node_index << std::endl;
+  //   } else {
+  //     std::cout << "dom tree node: nullptr" << std::endl;
+  //   }
+  //   auto root_group = groups_[i]->FindRoot();
+  //   // find root_group index in groups_
+  //   size_t root_group_index = 0xffffffff;
+  //   for (size_t j = 0; j < groups_.size(); j++) {
+  //     if (groups_[j] == root_group) {
+  //       root_group_index = j;
+  //       break;
+  //     }
+  //   }
+  //   std::cout << "root nid: " << root_group_index << std::endl;
+  // }
+  // ----------------------------------------------------------------------------------
+  // zz: for manual fusion end
+  // ----------------------------------------------------------------------------------
+}
+
+// zz: for random fuse testing
+std::vector<GraphPartitioner::Group*> GraphPartitioner::GetGroups() {
   return std::move(groups_);
 }
 
@@ -156,9 +367,9 @@ bool GraphPartitioner::CheckPath(IndexedForwardGraph::Node* src, IndexedForwardG
 }
 
 OpPatternKind CombinePattern(OpPatternKind lhs, OpPatternKind rhs) {
-  if (lhs > relay::kBroadcast && rhs > relay::kBroadcast) {
-    LOG(FATAL) << "Cannot merge two complex group together";
-  }
+  // if (lhs > relay::kBroadcast && rhs > relay::kBroadcast) {
+  //   LOG(FATAL) << "Cannot merge two complex group together";
+  // } // zz: comment this to enable complex fuse first
   if (lhs > rhs) return lhs;
   return rhs;
 }
@@ -170,10 +381,11 @@ void GraphPartitioner::MergeFromTo(Group* child, Group* parent) {
   // update the number of nodes of the parent group
   parent->num_nodes += child->num_nodes;
   parent->args_num += child->args_num;
+  parent->num_compute_nodes += child->num_compute_nodes; // zz: add num_compute_nodes
   child->parent = parent;
   // update anchor ref and pattern
   if (child->anchor_ref != nullptr) {
-    ICHECK(parent->anchor_ref == nullptr);
+    // ICHECK(parent->anchor_ref == nullptr); // zz: comment this to enable complex fuse first
     parent->anchor_ref = child->anchor_ref;
     parent->pattern = CombinePattern(child->pattern, parent->pattern);
   }
@@ -233,6 +445,7 @@ size_t GraphPartitioner::CountAdditionalArgs_(const TensorTypeNode* ttype, bool
     }
   }
   if (with_strides && any_dims > 0) any_dims += ttype->shape.size();
+  // std::cout << "[ZZDEBUG] GraphPartitioner::CountAdditionalArgs_: any_dims = (all 0 in example) " << any_dims << std::endl;
   return any_dims;
 }
 
@@ -344,9 +557,24 @@ void GraphPartitioner::InitGroups(const IndexedForwardGraph& graph) {
     }
     group_node->args_num = args_counter(graph_node->ref);
     groups_[nid] = group_node;
+
+    // zz: check whether the node is a compute node
+    group_node->num_compute_nodes = 0;
+    if (auto call_node = GetRef<ObjectRef>(graph_node->ref).as<CallNode>()) {
+      // if (auto call_op = call_node->op.as<OpNode>()) {
+        // check if call_op is conv2d, dense, max_pool2d
+        // if (call_op->name == "nn.conv2d" || call_op->name == "nn.dense" || call_op->name == "nn.max_pool2d") {
+      if (isComputeOp(call_node)) {
+        group_node->num_compute_nodes = 1;
+      }
+      // }
+    }
+    // zz: set num_fop_nodes
+    group_node->num_fop_nodes = 0;
   }
 }
 
+// zz: debug log added
 void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
                                const DominatorTree& post_dom_tree,  //
                                int phase) {
@@ -356,6 +584,15 @@ void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
     auto* dom_node = post_dom_tree.nodes[nid];
     Group* group_node = groups_[nid];
     ICHECK(group_node != nullptr);
+    // std::cout << "================================================================================================================================================================================" << std::endl;
+    // std::cout << "[ZZDEBUG] RunFuse phase " << std::dec << phase << " nid " << nid << ":" << std::endl;
+    // std::cout << "  (zz) try to fuse node: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph_node->ref))) << std::endl;
+    // if (dom_node->parent != nullptr) {
+    //   std::cout << "  (zz) to node:" << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(dom_node->parent->gnode->ref))) << std::endl;
+    // }
+    // else {
+    //   std::cout << "  (zz) to node: nullptr" << std::endl;
+    // }
     postpone_node_ = nullptr;
     // Check if the fusing of some inputs was postponed
     if (postponed_fusing_map_.count(graph_node)) {
@@ -366,26 +603,47 @@ void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
           auto* src = it->second;
           auto* snode = post_dom_tree.nodes[src->index]->parent->gnode;
           if (groups_[snode->index]->anchor_ref != nullptr) continue;
+          // std::cout << "  (zz) CommitFuse CountArgs_ (not in example): src= " << std::hex << (uintptr_t)src << ", dst = " << (uintptr_t)snode << std::endl;
+          // std::cout << "  [fuse] postponed_fusing_map_: " << std::endl;
+          // std::cout << "    zz: fuse node: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(src->ref))) << std::endl;
+          // std::cout << "    zz: to node: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(snode->ref))) << std::endl;
           CommitFuse(src, snode);
         }
       }
       postponed_fusing_map_.erase(graph_node);
     }
+    // std::cout << "  ----- (zz) after postponed_fusing_map_" << std::endl;
     // no actions for opaque nodes
-    if (group_node->pattern == kOpaque) continue;
+    // if (group_node->pattern == kOpaque) continue;
+    if (group_node->pattern == kOpaque) {
+      // std::cout << "  [no fuse] opaque node" << std::endl;
+      continue;
+    }
     // no actions needed if the current node have no dominator
-    if (dom_node->parent == nullptr) continue;
+    // if (dom_node->parent == nullptr) continue;
+    if (dom_node->parent == nullptr) {
+      // std::cout << "  [no fuse] no dominator" << std::endl;
+      continue;
+    }
     ICHECK(!graph_node->extern_ref);
     size_t dom_parent_gindex = dom_node->parent->gnode->index;
 
     // refuse the fusion if too many ops are going to be fused together
-    if (CountFusedNodesWithNewChild(graph_node, dom_node->parent->gnode) > max_fuse_depth_)
+    // if (CountFusedNodesWithNewChild(graph_node, dom_node->parent->gnode) > max_fuse_depth_)
+    //   continue;
+    if (CountFusedNodesWithNewChild(graph_node, dom_node->parent->gnode) > max_fuse_depth_) {
+      // std::cout << "  [no fuse] CountFusedNodesWithNewChild exceed" << std::endl;
       continue;
+    }
     // Refuse the fusion if too many arguments are going to be in the fused function
     if (max_function_args_ > 0) {
       auto limit = CountArgsLimit_(graph_node);
       if (limit > 0) {
+        // if (CountFusedArgs(graph, graph_node) > limit) {
+        //   continue;
+        // }
         if (CountFusedArgs(graph, graph_node) > limit) {
+          // std::cout << "  [no fuse] CountFusedArgs exceed" << std::endl;
           continue;
         }
       }
@@ -393,33 +651,58 @@ void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
 
     if (phase == 2) {
       // Fuse injective ops into intermediate tuples, if any
-      if (group_node->pattern > relay::kInjective) continue;
+      // if (group_node->pattern > relay::kInjective) continue;
+      if (group_node->pattern > relay::kInjective) {
+        // std::cout << "  [no fuse] phase2 pattern > kInjective (need to be finished in phase 0/1)" << std::endl;
+        continue;
+      }
       Group* dom_parent_group = groups_[dom_parent_gindex];
       Group* dom_root_group = dom_parent_group->FindRoot();
       // If dom node group has a tuple as its root, we do not fuse tuple fields into it
-      if (dom_root_group->pattern == relay::kTuple) continue;
+      // if (dom_root_group->pattern == relay::kTuple) continue;
+      if (dom_root_group->pattern == relay::kTuple) {
+        // std::cout << "  [no fuse] phase2 dom_root_group->pattern == kTuple" << std::endl;
+        continue;
+      }
       if (dom_parent_group->pattern == kTuple && dom_root_group->pattern <= relay::kInjective) {
         // Now we know the tuple has been fused into subsequent injective ops
         auto fcond = [](OpPatternKind kind, bool is_sink) { return kind <= kInjective; };
         // dom_root_group can also be tuple, as in inception layers
         // CheckPath is needed to avoid fusing two intermediate tuples
         if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
+          // std::cout << "  (zz) CommitFuse phase == 2 (not in example): src= " << std::hex << (uintptr_t)graph_node << ", dst = " << (uintptr_t)(dom_node->parent->gnode) << std::endl;
+          // std::cout << "  [fuse] phase 2 " << std::endl;
           CommitFuse(graph_node, dom_node->parent->gnode);
         }
+        // else {
+        //   std::cout << "  [no fuse] phase2 CheckPath failed" << std::endl;
+        // }
       }
+      // std::cout << "  [no fuse] phase2 not fuse case" << std::endl;
       continue;
     }
+    // std::cout << "  ----- (zz) after phase == 2" << std::endl;
 
     // Skip if current node is already fused to the parent.
     if (groups_[dom_parent_gindex] != nullptr &&
         group_node->FindRoot() == groups_[dom_parent_gindex]->FindRoot()) {
+      // std::cout << "  [no fuse] already fused to parent" << std::endl;
       continue;
     }
     // Do not fuse into tuple for now
-    if (groups_[dom_parent_gindex]->pattern == kTuple) continue;
+    // if (groups_[dom_parent_gindex]->pattern == kTuple) continue;
+    if (groups_[dom_parent_gindex]->pattern == kTuple) {
+      // std::cout << "  [no fuse] dom kTuple" << std::endl;
+      continue;
+    }
     // Try to fuse current node to its post-dominator.
     if (group_node->pattern == kOutEWiseFusable) {
-      if (phase != 0) continue;
+      // std::cout << "  ----- (zz) kOutEWiseFusable" << std::endl;
+      // if (phase != 0) continue;
+      if (phase != 0) {
+        // std::cout << "  [no fuse] src kOutEWiseFusable, phase != 0" << std::endl;
+        continue;
+      }
       // Path for OutEWiseFusable: conv2d
       // Check if the dominator relation is elemwise.
       if (dom_node->parent != nullptr && dom_node->pattern == kElemWise) {
@@ -427,10 +710,19 @@ void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
         // The fuse can be executed if all the intermediate ops are still broadcast.
         auto fcond = [](OpPatternKind kind, bool is_sink) { return kind <= kBroadcast; };
         if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
+          // std::cout << "  [fuse] src kOutEWiseFusable" << std::endl;
           CommitFuse(graph_node, dom_node->parent->gnode);
+          // std::cout << "        after CommitFuse: dom_node->parent->gnode = " << std::hex << (uintptr_t)(dom_node->parent->gnode) << ", dom_node->pattern = " << std::dec << dom_node->pattern << std::endl;
         }
+        // else {
+        //   std::cout << "  [no fuse] src kOutEWiseFusable CheckPath failed" << std::endl;
+        // }
       }
+      // else {
+      //   std::cout << "  [no fuse] src kOutEWiseFusable, dom !kElemWise" << std::endl;
+      // }
     } else if (group_node->pattern <= kBroadcast) {
+      // std::cout << "  ----- (zz) <= kBroadcast" << std::endl;
       // Pre-condition: can only be fused to parent which is injective or reduction.
       if (dom_node->parent != nullptr &&
           (dom_node->pattern <= kInjective || dom_node->pattern == kCommReduce)) {
@@ -447,24 +739,189 @@ void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
           }
         };
         if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
+          // std::cout << "  [fuse] src <= kBroadcast" << std::endl;
           CommitFuse(graph_node, dom_node->parent->gnode);
         }
       }
     } else if (group_node->pattern == kInjective || group_node->pattern == kTuple) {
+      // std::cout << "  ----- (zz) kInjective - kTuple" << std::endl;
+      // std::cout << "  (zz) try fuse to dom_parent: group_node->pattern == kInjective || kTuple" << std::endl;
       // defer injective fusion to second phase.
       // so conv2d always finishes fusing.
-      if (phase != 1) continue;
+      // if (phase != 1) continue;
+      if (phase != 1) {
+        // std::cout << "  [no fuse] src kInjective || kTuple, phase != 1" << std::endl;
+        continue;
+      }
       // Check if all path are injective.
       auto fcond = [](OpPatternKind kind, bool is_sink) { return kind <= kInjective; };
       if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
+        // std::cout << "    (zz) CommitFuse kInjective || kInjective: src= " << std::hex << (uintptr_t)graph_node << ", dst = " << (uintptr_t)(dom_node->parent->gnode) << std::endl;
+        // std::cout << "  [fuse] src kInjective || kTuple" << std::endl;
         CommitFuse(graph_node, dom_node->parent->gnode);
       }
+      else {
+        // std::cout << "  [no fuse] src kInjective || kTuple CheckPath failed" << std::endl;
+      }
     } else {
       // do nothing.
       ICHECK(group_node->pattern == kCommReduce);
+      // std::cout << "  [no fuse] src kCommReduce" << std::endl;
     }
   }
 }
 
+// zz: path checker for flexible fusion
+bool GraphPartitioner::CheckPathObfFlex(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink) {
+  // src and sink must be call node
+  if (!src->ref->IsInstance<CallNode>() || !sink->ref->IsInstance<CallNode>()) {
+    // std::cout << "[ZZDEBUG] CheckPathObfFlex: src or sink not CallNode" << std::endl;
+    return false;
+  }
+  // sink call not arg only contain one call node input (which is src)
+  auto sink_call = Downcast<Call>(GetRef<Call>(static_cast<const CallNode*>(sink->ref)));
+  size_t sink_input_size = 0;
+  for (auto arg : sink_call->args) {
+    if (arg->IsInstance<CallNode>()) {
+      sink_input_size++;
+    }
+  }
+  if (sink_input_size != 1) {
+    // std::cout << "[ZZDEBUG] CheckPathObfFlex: sink_input_size != 1" << std::endl;
+    return false;
+  }
+
+  // src outpus size must be 1, and the output is sink
+  size_t src_output_size = 0;
+  for (auto link = src->outputs.head; link != nullptr; link = link->next) {
+    src_output_size++;
+  }
+  if (src_output_size != 1) {
+    return false;
+  }
+  if (src->outputs.head->value.node != sink) {
+    return false;
+  }
+  return true;
+}
+
+// zz: compute op checker for flexible fusion
+size_t GraphPartitioner::CountComputeOp(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink) {
+  // find root
+  Group* src_grp = groups_[src->index];
+  Group* sink_grp = groups_[sink->index];
+  Group* src_root = src_grp->FindRoot();
+  Group* sink_root = sink_grp->FindRoot();
+  // check if src and sink are in the same group
+  if (src_root == sink_root) {
+    return src_root->num_compute_nodes;
+  }
+  return src_root->num_compute_nodes + sink_root->num_compute_nodes;
+}
+
+// zz: flexible fusion for obfuscation
+void GraphPartitioner::RunFuseObfFlex(const IndexedForwardGraph& graph, const DominatorTree& post_dom_tree) {
+  for (size_t nid = 0; nid < groups_.size(); ++nid) {
+    // the group of current node has been specified already.
+    auto* graph_node = graph.post_dfs_order[nid];
+    auto* dom_node = post_dom_tree.nodes[nid];
+    Group* group_node = groups_[nid];
+
+    // std::cout << "================================================================================================================================================================================" << std::endl;
+    // std::cout << "[ZZDEBUG] RunFuseObfFlex " << std::dec << " nid " << nid << ":" << std::endl;
+    // std::cout << "  (zz) try to fuse node: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph_node->ref))) << std::endl;
+    // if (dom_node->parent != nullptr) {
+    //   std::cout << "  (zz) to node:" << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(dom_node->parent->gnode->ref))) << std::endl;
+    // }
+    // else {
+    //   std::cout << "  (zz) to node: nullptr" << std::endl;
+    // }
+
+    // skip if no dom parent
+    if (dom_node->parent == nullptr) {
+      // std::cout << "  [RunFuseObfFlex, no fuse] no dominator" << std::endl;
+      continue;
+    }
+
+    // skip if current node is already fused to the parent.
+    size_t dom_parent_gindex = dom_node->parent->gnode->index;
+    if (groups_[dom_parent_gindex] != nullptr &&
+        group_node->FindRoot() == groups_[dom_parent_gindex]->FindRoot()) {
+      // std::cout << "  [RunFuseObfFlex, no fuse] already fused to parent" << std::endl;
+      continue;
+    }
+
+    // try to fuse
+    if (CheckPathObfFlex(graph_node, dom_node->parent->gnode)) {
+      if (CountComputeOp(graph_node, dom_node->parent->gnode) <= max_compute_nodes_) {
+        // std::cout << "  [RunFuseObfFlex, fuse] CountComputeOp <= max_compute_nodes_" << std::endl;
+        CommitFuse(graph_node, dom_node->parent->gnode);
+      }
+      // else {
+      //   std::cout << "  [RunFuseObfFlex, no fuse] CountComputeOp > max_compute_nodes_" << std::endl;
+      // }
+    }
+  }
+}
+
+
+// zz: determine fops
+void GraphPartitioner::SelectFops(const IndexedForwardGraph& graph, const DominatorTree& post_dom_tree, int phase) {
+  for (size_t nid = 0; nid < groups_.size(); ++nid) {
+    auto* graph_node = graph.post_dfs_order[nid];
+    // auto* dom_node = post_dom_tree.nodes[nid];
+    Group* group_node = groups_[nid];
+    // std::cout << "================================================================================================================================================================================" << std::endl;
+    // std::cout << "[ZZDEBUG] SelectFops phase " << std::dec << phase << " nid " << nid << ":" << std::endl;
+    // std::cout << "  (zz) try to select fops: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph_node->ref))) << std::endl;
+
+    auto root_group = group_node->FindRoot();
+    if (root_group->num_fop_nodes >= max_fop_depth_) {
+      // std::cout << "  [SelectFops, no select] num_fop_nodes >= max_fop_depth_ (" << std::dec << root_group->num_fop_nodes << " vs. " << max_fop_depth_ << ")" << std::endl;
+      continue;
+    }
+
+    // add compute nodes
+    if (auto call_node = GetRef<ObjectRef>(graph_node->ref).as<CallNode>()) {
+      // if (auto call_op = call_node->op.as<OpNode>()) {
+      if (fake_ops_.count(graph_node->ref)) {
+        // std::cout << "  [SelectFops, no select] already in fake_ops_" << std::endl;
+        continue;
+      }
+
+      if (isComputeOp(call_node)) {
+        root_group->num_fop_nodes++;
+        fake_ops_.insert(graph_node->ref);
+        // std::cout << "  [SelectFops, select] compute node: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph_node->ref))) << std::endl;
+      }
+
+      if (phase > 0 && isSmallComputeOp(call_node)) {
+        root_group->num_fop_nodes++;
+        fake_ops_.insert(graph_node->ref);
+        // std::cout << "  [SelectFops, select] small compute node: " << std::endl << PrettyPrint(GetRef<Expr>(static_cast<const ExprNode*>(graph_node->ref))) << std::endl;
+      }
+      // }
+
+      // std::cout << "  [SelectFops, no select] not compute node" << std::endl;
+    }
+    // else {
+    //   std::cout << "  [SelectFops, no select] not call node" << std::endl;
+    // }
+  }
+}
+
+std::unordered_set<const Object*> GraphPartitioner::GetFops() {
+  return std::move(fake_ops_);
+}
+
+// zz: for random fuse testing
+void GraphPartitioner::MannualFuse(const IndexedForwardGraph& graph, int src_inx) {
+  auto post_dom_tree = DominatorTree::PostDom(arena_, graph);
+  auto* graph_node = graph.post_dfs_order[src_inx];
+  auto* dom_node = post_dom_tree.nodes[src_inx];
+  // Group* group_node = groups_[src_inx];
+  CommitFuse(graph_node, dom_node->parent->gnode);
+}
+
 }  // namespace relay
 }  // namespace tvm
diff --git a/src/relay/analysis/graph_partitioner.h b/src/relay/analysis/graph_partitioner.h
index b3f934b97..8bbbd51a8 100644
--- a/src/relay/analysis/graph_partitioner.h
+++ b/src/relay/analysis/graph_partitioner.h
@@ -163,11 +163,17 @@ class DominatorTree {
 class GraphPartitioner {
  public:
   explicit GraphPartitioner(support::Arena* arena, int opt_level, size_t max_fuse_depth,
-                            size_t max_function_args)
+                            // size_t max_function_args)
+                            size_t max_function_args,
+                            size_t max_compute_nodes=3, size_t max_fop_depth=3) // zz: add max_compute_nodes, max_fop_depth
+                            // size_t max_compute_nodes=2, size_t max_fop_depth=3) // zz: max_compute_nodes 2 for transformers to avoid bug (todo: make this configurable)
       : arena_(arena),
         opt_level_(opt_level),
         max_fuse_depth_(max_fuse_depth),
-        max_function_args_(max_function_args) {}
+        // max_function_args_(max_function_args) {}
+        max_function_args_(max_function_args),
+        max_compute_nodes_(max_compute_nodes),  // zz: add max_compute_nodes
+        max_fop_depth_(max_fop_depth) {} // zz: add max_fop_depth
   /*!
    * \brief Group as a union find data structure.
    */
@@ -187,6 +193,10 @@ class GraphPartitioner {
      * \brief The number of nodes belonging to this group
      */
     uint32_t num_nodes{1};
+    // zz: add num_compute_nodes for flexible op fusion
+    uint32_t num_compute_nodes{0};
+    // zz: add num_fop_nodes for flexible op fusion
+    uint32_t num_fop_nodes{1};
     /*!
      * \brief The number of function arguments belonging to this group
      */
@@ -206,6 +216,13 @@ class GraphPartitioner {
    */
   std::vector<Group*> Partition(const IndexedForwardGraph& graph);
 
+  // zz: for random fuse testing
+  void PartitionWithoutMov(const IndexedForwardGraph& graph);
+  std::vector<Group*> GetGroups();
+  void MannualFuse(const IndexedForwardGraph& graph, int src_inx);
+  // zz: for fake op selection
+  std::unordered_set<const Object*> GetFops();
+
  private:
   /*! \brief The internal arena for temporary space. */
   support::Arena* arena_;
@@ -215,8 +232,14 @@ class GraphPartitioner {
   size_t max_fuse_depth_;
   /*! \brief The maximum number of arguments in one fused function */
   size_t max_function_args_;
+  // zz: add max_compute_nodes_ for flexible op fusion
+  size_t max_compute_nodes_;
+  // zz: add max_fop_depth_ for flexible op fusion
+  size_t max_fop_depth_;
   /*! \brief The internal groups. */
   std::vector<Group*> groups_;
+  // zz: op lists for fake op insertion
+  std::unordered_set<const Object*> fake_ops_;
   /*! \brief internal field used for deduplication */
   std::unordered_set<IndexedForwardGraph::Node*> visited_;
   /*! \brief The map with nodes which were postponed for fusing. */
@@ -246,6 +269,8 @@ class GraphPartitioner {
    */
   template <typename F>
   bool CheckPath(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink, F fcond);
+  // zz: path checker for flexible fusion
+  bool CheckPathObfFlex(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink);
 
   /*!
    * \brief Merge the child group to the parent.
@@ -266,6 +291,8 @@ class GraphPartitioner {
   void CommitFuse(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink);
 
   size_t CountNodesUptoSink_(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink);
+  // zz: compute op checker for flexible fusion
+  size_t CountComputeOp(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink);
   // Count the number of additional arguments. In the case of dynamic shape,
   // generated function takes several additional arguments, such as the sizes of
   // the dynamic dimensions and strides.
@@ -302,6 +329,11 @@ class GraphPartitioner {
 
   // execute the fusion algorithm.
   void RunFuse(const IndexedForwardGraph& graph, const DominatorTree& post_dom_tree, int phase);
+  // zz: flexible fusion for obfuscation
+  void RunFuseObfFlex(const IndexedForwardGraph& graph, const DominatorTree& post_dom_tree);
+
+  // zz: determine fops
+  void SelectFops(const IndexedForwardGraph& graph, const DominatorTree& post_dom_tree, int phase);
 };
 
 }  // namespace relay
diff --git a/src/relay/backend/annotate_used_memory.cc b/src/relay/backend/annotate_used_memory.cc
index 8e7ab68cb..150952c92 100644
--- a/src/relay/backend/annotate_used_memory.cc
+++ b/src/relay/backend/annotate_used_memory.cc
@@ -229,6 +229,31 @@ Pass AnnotateUsedMemory() {
   return CreateModulePass(pass_func, 0, "AnnotateUsedMemory", {"ToANormalForm", "InferType"});
 }
 
+// zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+Pass AnnotateUsedMemoryWithMainFuncName(String main_func_name) {
+  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func = [=](IRModule mod,
+                                                                            PassContext ctx) {
+    // std::cout << "[ZZDEBUG] AnnotateUsedMemoryWithMainFuncName input mod:" << std::endl << PrettyPrint(mod) << "[ZZDEBUG] AnnotateUsedMemoryWithMainFuncName input mod end." << std::endl;
+    GlobalVar gv = mod->GetGlobalVar(main_func_name);
+    Function main_func = Downcast<Function>(mod->Lookup(main_func_name));
+
+    // Perform liveness analysis to determine what tensors are 'live' at each functions callsite.
+    support::Arena arena;
+    ControlFlowGraph cfg = ControlFlowGraph::Create(&arena, main_func);
+    UseDefAnalysis use_def = UseDefAnalysis::Analyze(cfg);
+    LivenessAnalysis lva = LivenessAnalysis::Analyze(cfg, use_def);
+
+    auto new_main_func = backend::AnnotateUsedMemoryMutator(mod, cfg, lva)(main_func);
+    // std::cout << "[ZZDEBUG] AnnotateUsedMemoryWithMainFuncName mod before update main:" << std::endl << PrettyPrint(mod) << "[ZZDEBUG] AnnotateUsedMemoryWithMainFuncName mod before update main end." << std::endl;
+    if (!new_main_func.same_as(main_func)) {
+      mod->Update(gv, new_main_func);
+    }
+    // std::cout << "[ZZDEBUG] AnnotateUsedMemoryWithMainFuncName mod after update main:" << std::endl << PrettyPrint(mod) << "[ZZDEBUG] AnnotateUsedMemoryWithMainFuncName mod after update main end." << std::endl;
+    return mod;
+  };
+  return CreateModulePass(pass_func, 0, "AnnotateUsedMemory", {"ToANormalForm", "InferType"});
+}
+
 TVM_REGISTER_GLOBAL("relay._transform.AnnotateUsedMemory").set_body_typed(AnnotateUsedMemory);
 
 }  // namespace transform
diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc
index f698c654d..49fad563b 100644
--- a/src/relay/backend/aot_executor_codegen.cc
+++ b/src/relay/backend/aot_executor_codegen.cc
@@ -61,6 +61,79 @@ namespace backend {
 using StorageMap =
     std::unordered_map<Expr, StorageInfo, runtime::ObjectPtrHash, runtime::ObjectPtrEqual>;
 
+// zz: visitor for extract non prim obfuscated calls in main function
+class NonPrimCallCollector : public ExprVisitor {
+ public:
+  NonPrimCallCollector() {}
+
+  void VisitExpr_(const FunctionNode* func_node) final {
+    // if (const auto* func_node = call_node->op.as<FunctionNode>()) {
+    if (!func_node->attrs.HasNonzeroAttr(attr::kPrimitive)) {
+      Function func = GetRef<Function>(func_node);
+      non_prim_call_func_.push_back(func); // skip visiting inside function
+    }
+    else {
+      ExprVisitor::VisitExpr_(func_node);
+    }
+  }
+
+
+  std::vector<Function> GetNonPrimCallFuncs() const { return non_prim_call_func_; }
+
+ private:
+  std::vector<Function> non_prim_call_func_;
+};
+// zz: mutaotr to add kObf attribute to non prim functions
+class NonPrimCallAnnotator : public ExprMutator {
+ public:
+  NonPrimCallAnnotator(std::vector<Function> non_prim_call_func) : non_prim_call_func_{non_prim_call_func} {}
+
+  Expr VisitExpr_(const FunctionNode* func_node) final {
+    size_t func_index = 0;
+    // find if the function collected
+    for (func_index = 0; func_index < non_prim_call_func_.size(); func_index++) {
+      if (func_node == non_prim_call_func_[func_index].get()) {
+        break;
+      }
+    }
+    if (func_index < non_prim_call_func_.size()) {
+      Function func = GetRef<Function>(func_node);
+      Function func_obf =  WithAttrs(func, {{attr::kObf, Integer(func_index)}}); // annoate with obf function index
+      non_prim_obfid_call_func_.push_back(std::make_pair(func_index, func));
+      return func_obf; // skip visiting inside function
+    }
+    else {
+      return ExprMutator::VisitExpr_(func_node);
+    }
+  }
+
+  std::vector<std::pair<int, Function> > GetNonPrimCallFuncs() const { return non_prim_obfid_call_func_; }
+
+ private:
+  std::vector<Function> non_prim_call_func_;
+  std::vector<std::pair<int, Function> > non_prim_obfid_call_func_;
+};
+
+using SubModItem = std::pair<String, IRModule>; // zz: obf_mod_name (obf main name), obf_mod
+
+// // zz: visitor for extract let bindings
+// class LetBindingCollector : public ExprVisitor {
+//  public:
+//   LetBindingCollector() {}
+
+//   void VisitExpr_(const LetNode* let_node) final {
+//     std::cout << "[ZZDEBUG] LetBindingCollector VisitExpr_ LetNode: " << std::endl << PrettyPrint(GetRef<Expr>(let_node)) << std::endl;
+//     let_bindings_.insert({let_node->var, GetRef<Expr>(let_node)});
+//     ExprVisitor::VisitExpr_(let_node);
+//   }
+
+//   std::unordered_map<Expr, Expr, runtime::ObjectPtrHash, runtime::ObjectPtrEqual> GetLetBindings() const { return let_bindings_; }
+
+//  private:
+//   // std::vector<const LetNode*> let_bindings_;
+//   std::unordered_map<Expr, Expr, runtime::ObjectPtrHash, runtime::ObjectPtrEqual> let_bindings_;
+// };
+
 /**
  * This is an on demand allocator for AOT. A new temporary
  * (storage allocator identifier) is allocated for each operation.
@@ -68,6 +141,8 @@ using StorageMap =
 class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
  public:
   AOTOnDemandAllocator() : transform::DeviceAwareExprVisitor(Optional<IRModule>()) {}
+  // zz: add let_bindings
+  // AOTOnDemandAllocator(std::unordered_map<Expr, Expr, runtime::ObjectPtrHash, runtime::ObjectPtrEqual> let_bindings = {}) : transform::DeviceAwareExprVisitor(Optional<IRModule>()), let_bindings_{let_bindings} {}
 
   // run the visitor on a global function.
   void Run(const Function& func) { VisitExpr(func); }
@@ -167,7 +242,74 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
     AssignReturnSid(expr);
   }
 
-  void VisitExpr_(const IfNode* op) final { LOG(FATAL) << "if is not supported."; }
+  // zz: for debugging storage_device_map_
+  void PrintStorageDeviceMap() {
+    std::cout << "[ZZDEBUG] current storage_device_map_:" << std::endl;
+    for (auto& kv : storage_device_map_) {
+      std::cout << "  (zz) key: " << kv.first << std::endl;
+      std::cout << "  (zz) value: " << kv.second << std::endl;
+    }
+  }
+
+  // void VisitExpr_(const IfNode* op) final { LOG(FATAL) << "if is not supported."; }
+  // zz: hanlde IfNode
+  void VisitExpr_(const IfNode* op) final {
+    if (!if_expr_stack_.empty()) {
+      // propagate storage from parent if to current if
+      auto parent_if_expr = if_expr_stack_.back();
+      auto parent_if = Downcast<If>(parent_if_expr);
+      auto parent_if_true_expr = parent_if->true_branch;
+      parent_if_true_expr = IgnoreLetInIfBranch(parent_if_true_expr);
+      parent_if_true_expr = IgnoreOnDevice(parent_if_true_expr);
+      auto parent_if_false_expr = parent_if->false_branch;
+      parent_if_false_expr = IgnoreLetInIfBranch(parent_if_false_expr);
+      parent_if_false_expr = IgnoreOnDevice(parent_if_false_expr);
+
+      // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode parent_if_true_expr: " << std::endl << PrettyPrint(parent_if_true_expr) << std::endl;
+      // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode parent_if_false_expr: " << std::endl << PrettyPrint(parent_if_false_expr) << std::endl;
+
+      auto cur_if = GetRef<Expr>(op);
+      // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode cur_if: " << std::endl << PrettyPrint(cur_if) << std::endl;
+      if (cur_if.same_as(parent_if_true_expr) || cur_if.same_as(parent_if_false_expr)) {
+        // zz: if expr is the true or false branch of parent if, then it should have the same storage as the parent if
+        // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode: found exsiting storage" << std::endl;
+        ICHECK(storage_device_map_.find(parent_if_expr) != storage_device_map_.end()) << "parent if storage not existing";
+        auto si = storage_device_map_[parent_if_expr];
+        // check storage size
+        int s_idx = 0;
+        for (const auto& ttype : FlattenTupleType(op->checked_type())) { 
+          ICHECK_EQ(GetMemorySizeBytes(ttype), si->storage_sizes_in_bytes[s_idx++]);
+        }
+        storage_device_map_[GetRef<Expr>(op)] = si;
+      }
+    }
+
+    // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode: " << std::endl << PrettyPrint(GetRef<Expr>(op)) << std::endl;
+    // PrintStorageDeviceMap();
+    VisitExpr(op->cond);
+    // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode: cond visited: " << std::endl << PrettyPrint(op->cond) << std::endl;
+    // PrintStorageDeviceMap();
+    if_expr_stack_.push_back(GetRef<Expr>(op));
+    VisitExpr(op->true_branch);
+    // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode: true_branch visited: " << std::endl << PrettyPrint(op->true_branch) << std::endl;
+    // PrintStorageDeviceMap();
+    VisitExpr(op->false_branch);
+    // std::cout << "[ZZDEBUG] AOTOnDemandAllocator VisitExpr_ IfNode: false_branch visited: " << std::endl << PrettyPrint(op->false_branch) << std::endl;
+    // PrintStorageDeviceMap();
+    if_expr_stack_.pop_back();
+
+    // zz: set storage for if node (should be set alreay in CreateStorage?)
+    auto true_expr = op->true_branch;
+    true_expr = IgnoreLetInIfBranch(true_expr); // zz: ignore let to get true expr in ifnode branches
+    auto true_storage = GetStorage(true_expr);
+    auto false_expr = op->false_branch;
+    false_expr = IgnoreLetInIfBranch(false_expr); // zz: ignore let to get false expr in ifnode branches
+    auto false_storage = GetStorage(false_expr);
+    ICHECK_EQ(true_storage.same_as(false_storage), true) << "if node true and false branch should have the same storage" << std::endl
+                                                          << "true branch:" << std::endl << PrettyPrint(true_expr) << std::endl
+                                                          << "false branch:" << std::endl << PrettyPrint(false_expr) << std::endl;
+    storage_device_map_[GetRef<Expr>(op)] = true_storage;
+  }
 
   void PreVisitLetBinding_(const Var& var, const Expr& value) final {
     VisitExpr(value);
@@ -175,6 +317,13 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
     storage_device_map_[var] = si;
   }
 
+  // zz: Overriding PostVisitLet_ is needed when there are if nodes
+  void PostVisitLet_(const LetNode* let_node) final {
+    // zz: set storage for let node with si for var
+    StorageInfo si = GetStorage(let_node->var);
+    storage_device_map_[GetRef<Expr>(let_node)] = si;
+  }
+
  private:
   void AssignReturnSid(Expr e) {
     if (storage_device_map_.find(e) != storage_device_map_.end()) {
@@ -228,6 +377,46 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
     return it->second;
   }
 
+  // zz: [helper function] ignore let to get true expr in ifnode branches (ref: src/relay/op/memory/on_device.h IgnoreOnDevice)
+  // Expr IgnoreLet(const Expr& expr) {
+  Expr IgnoreLetInIfBranch(const Expr& expr) {
+    // if (const LetNode* let = expr.as<LetNode>()) {
+    //   return IgnoreLet(let->value);
+    // }
+    // return expr;
+
+    // zz: if branch body is the first let, need to get the final let
+    // if (const VarNode* var = expr.as<VarNode>()) {
+    //   std::cout << "[ZZDEBUG] IgnoreLetInIfBranch expr is var: " << std::endl << PrettyPrint(expr) << std::endl;
+    //   auto it = let_bindings_.find(GetRef<Expr>(var));
+    //   if (it != let_bindings_.end()) {
+    //     std::cout << "[ZZDEBUG] IgnoreLetInIfBranch var related let: " << std::endl << PrettyPrint(it->second) << std::endl;
+    //     return IgnoreLetInIfBranch(it->second);
+    //   }
+    // }
+    if (const LetNode* let = expr.as<LetNode>()) {
+      while (const LetNode* let_body = let->body.as<LetNode>()) {
+        // std::cout << "[ZZDEBUG] GetLetBody: let var: " << std::endl << PrettyPrint(let->var) << std::endl;
+        if (let_body != let) {
+          let = let_body;
+        }
+        else {
+          break;
+        }
+      }
+      // std::cout << "[ZZDEBUG] GetLetBody: let var: " << std::endl << PrettyPrint(let->var) << std::endl;
+      return let->value;
+    }
+    // else {
+    //   std::cout << "[ZZDEBUG] IgnoreLetInIfBranch expr is not let: " << std::endl << PrettyPrint(expr) << std::endl;
+    //   // check var
+    //   if (const VarNode* var = expr.as<VarNode>()) {
+    //     std::cout << "[ZZDEBUG] IgnoreLetInIfBranch expr is var" << std::endl;
+    //   }
+    // }
+    return expr;
+  }
+
   /*!
    * \brief Create storage for the expression.
    */
@@ -240,6 +429,53 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
    * \brief Create storage to hold the result of evaluating \p expr in \p virtual_device.
    */
   void CreateStorage(const Expr& expr, const VirtualDevice& virtual_device) {
+    // zz: check whether expr is the true or false branch of current if scope
+    if (!if_expr_stack_.empty()) {
+      auto if_expr = if_expr_stack_.back();
+      // std::cout << "[ZZDEBUG] CreateStorage: if_expr_stack_ back: " << std::endl << PrettyPrint(if_expr) << std::endl;
+      auto donwncast_if_expr = Downcast<If>(if_expr);
+
+      auto if_true_expr = donwncast_if_expr->true_branch;
+      // std::cout << "[ZZDEBUG] CreateStorage: if_true_expr: " << std::endl << PrettyPrint(if_true_expr) << std::endl;
+      if_true_expr = IgnoreLetInIfBranch(if_true_expr); // zz: if body can be the frist let in the statement, and the let body is the final true expr
+      // std::cout << "[ZZDEBUG] CreateStorage: if_true_expr after GetLetBody: " << std::endl << PrettyPrint(if_true_expr) << std::endl;
+      if_true_expr = IgnoreOnDevice(if_true_expr);
+      // std::cout << "[ZZDEBUG] CreateStorage: if_true_expr after IgnoreOnDevice: " << std::endl << PrettyPrint(if_true_expr) << std::endl;
+
+      auto if_false_expr = donwncast_if_expr->false_branch;
+      // std::cout << "[ZZDEBUG] CreateStorage: if_false_expr: " << std::endl << PrettyPrint(if_false_expr) << std::endl;
+      if_false_expr = IgnoreLetInIfBranch(if_false_expr);
+      // std::cout << "[ZZDEBUG] CreateStorage: if_false_expr after GetLetBody: " << std::endl << PrettyPrint(if_false_expr) << std::endl;
+      if_false_expr = IgnoreOnDevice(if_false_expr);
+      // std::cout << "[ZZDEBUG] CreateStorage: if_false_expr after IgnoreOnDevice: " << std::endl << PrettyPrint(if_false_expr) << std::endl;
+
+      if (expr.same_as(if_true_expr) || expr.same_as(if_false_expr)) {
+        if(storage_device_map_.find(if_expr) != storage_device_map_.end()) {
+          // zz: if expr is the true or false branch of if, then it should have the same storage as the if expr
+          auto si = storage_device_map_[if_expr];
+          // check storage size
+          int s_idx = 0;
+          for (const auto& ttype : FlattenTupleType(expr->checked_type())) {
+            // std::cout << "[ZZDEBUG] CreateStorage: found exsiting storage" << std::endl;
+            // PrintStorageDeviceMap();
+            ICHECK_EQ(GetMemorySizeBytes(ttype), si->storage_sizes_in_bytes[s_idx++]) << "if expr storage size mismatch:" << std::endl << PrettyPrint(expr) << std::endl
+                                                                                      << "expr:" << std::endl << PrettyPrint(expr) << std::endl
+                                                                                      << "if expr:" << std::endl << PrettyPrint(if_expr) << std::endl
+                                                                                      << "true expr:" << std::endl << PrettyPrint(if_true_expr) << std::endl
+                                                                                      << "false expr:" << std::endl << PrettyPrint(if_false_expr) << std::endl
+                                                                                      << "si #ids: " << si->storage_ids.size() << std::endl << "si #sizes: " << si->storage_sizes_in_bytes.size() << std::endl
+                                                                                      << "ttype: " << ttype << std::endl
+                                                                                      << "GetMemorySizeBytes(ttype): " << GetMemorySizeBytes(ttype) << std::endl
+                                                                                      << "si: " << si << std::endl;
+          }
+          storage_device_map_[expr] = si;
+          return;
+        }
+      }
+    }
+
+    // ICHECK(storage_device_map_.find(expr) == storage_device_map_.end())
+    //     << "Storage map already contains expr " << PrettyPrint(expr);
     ICHECK(!virtual_device->IsFullyUnconstrained())
         << "invalid virtual device for expr:" << std::endl
         << PrettyPrint(expr);
@@ -253,6 +489,28 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
     }
     storage_device_map_[expr] = StorageInfo(std::move(storage_ids), std::move(virtual_devices),
                                             std::move(storage_sizes_in_bytes));
+
+    // zz: store storage info for current if scope
+    if (!if_expr_stack_.empty()) {
+      auto if_expr = if_expr_stack_.back();
+      // std::cout << "[ZZDEBUG] CreateStorage: try to store storage info for current if scope: " << std::endl << PrettyPrint(if_expr) << std::endl;
+      auto donwncast_if_expr = Downcast<If>(if_expr);
+      auto if_true_expr = donwncast_if_expr->true_branch;
+      auto if_false_expr = donwncast_if_expr->false_branch;
+      // zz: ignore let & on_device to get true expr in ifnode branches
+      if_true_expr = IgnoreLetInIfBranch(if_true_expr);
+      if_true_expr = IgnoreOnDevice(if_true_expr);
+      if_false_expr = IgnoreLetInIfBranch(if_false_expr);
+      if_false_expr = IgnoreOnDevice(if_false_expr);
+      // std::cout << "[ZZDEBUG] CreateStorage: if_true_expr for current if scope: " << std::endl << PrettyPrint(if_true_expr) << std::endl;
+      // std::cout << "[ZZDEBUG] CreateStorage: if_false_expr for current if scope: " << std::endl << PrettyPrint(if_false_expr) << std::endl;
+      // std::cout << "[ZZDEBUG] CreateStorage: current expr: " << std::endl << PrettyPrint(expr) << std::endl;
+
+      if (expr.same_as(if_true_expr) || expr.same_as(if_false_expr)) {
+        // std::cout << "[ZZDEBUG] CreateStorage: store storage info for current if scope: " << PrettyPrint(if_expr) << std::endl;
+        storage_device_map_[if_expr] = storage_device_map_[expr];
+      }
+    }
   }
 
   /*! \brief mapping of expression -> storageInfo */
@@ -263,6 +521,10 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
   std::vector<int> return_ids_;
   /*! \brief the data types of the return values */
   std::vector<TensorType> return_ttypes_;
+
+  // zz: for if handling
+  std::vector<Expr> if_expr_stack_;
+  // std::unordered_map<Expr, Expr, runtime::ObjectPtrHash, runtime::ObjectPtrEqual> let_bindings_;
 };
 
 /*! \brief Code generator for AOT executor */
@@ -521,7 +783,23 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     }
 
     tir::Stmt body = tir::SeqStmt::Flatten(func_call);
-    stmts_.push_back(body);
+    // stmts_.push_back(body);
+    // zz: handle if here: add flatten stmts to parent if scope
+    if (aot_codegen_if_expr_stack_.empty()) {
+      stmts_.push_back(body);
+    } else {
+      auto if_expr_scope = aot_codegen_if_expr_stack_.back();
+      if (if_expr_scope.second == 0) {
+        temp_stmts_[if_expr_scope.first].first.push_back(body);
+      }
+      else if (if_expr_scope.second == 1) {
+        temp_stmts_[if_expr_scope.first].second.push_back(body);
+      }
+      else {
+        ICHECK(false) << "if_expr_scope.second should be 0 or 1";
+      }
+
+    }
   }
 
   /*!
@@ -751,8 +1029,50 @@ class AOTExecutorCodegen : public MixedModeVisitor {
       LOG(FATAL) << "All OpNodes except for call_lowered should have been expanded";
     }
   }
+  // zz: add IfNode logic
   void VisitExpr_(const IfNode* op) override {
-    LOG(FATAL) << "All GlobalVarNodes should be removed before AOT executor's Codegen is called";
+    VisitExpr(op->cond);
+
+    aot_codegen_if_expr_stack_.push_back(std::make_pair(GetRef<Expr>(op), 0));
+    // std::cout << "[ZZDEBUG] AOTExecutorCodegen VisitExpr_ IfNode push true branch for: " << std::endl << PrettyPrint(GetRef<Expr>(op)) << std::endl;
+    // std::cout << "  (zz) current aot_codegen_if_expr_stack_ size: " << aot_codegen_if_expr_stack_.size() << std::endl;
+    VisitExpr(op->true_branch);
+    aot_codegen_if_expr_stack_.pop_back();
+    // std::cout << "[ZZDEBUG] AOTExecutorCodegen VisitExpr_ IfNode pop true branch for: " << std::endl << PrettyPrint(GetRef<Expr>(op)) << std::endl;
+    // std::cout << "  (zz) current aot_codegen_if_expr_stack_ size: " << aot_codegen_if_expr_stack_.size() << std::endl;
+
+    aot_codegen_if_expr_stack_.push_back(std::make_pair(GetRef<Expr>(op), 1));
+    // std::cout << "[ZZDEBUG] AOTExecutorCodegen VisitExpr_ IfNode push false branch for: " << std::endl << PrettyPrint(GetRef<Expr>(op)) << std::endl;
+    // std::cout << "  (zz) current aot_codegen_if_expr_stack_ size: " << aot_codegen_if_expr_stack_.size() << std::endl;
+    VisitExpr(op->false_branch);
+    aot_codegen_if_expr_stack_.pop_back();
+    // std::cout << "[ZZDEBUG] AOTExecutorCodegen VisitExpr_ IfNode pop false branch for: " << std::endl << PrettyPrint(GetRef<Expr>(op)) << std::endl;
+    // std::cout << "  (zz) current aot_codegen_if_expr_stack_ size: " << aot_codegen_if_expr_stack_.size() << std::endl;
+
+    // construct if stmt
+    auto cond = FindExpr(op->cond); // zz: FindExpr works well here
+    auto cond_buffer = tir::Buffer(cond[0], DataType::Int(8),
+                                  {tir::make_const(DataType::Int(32), 1)}, {}, 0, cond[0]->name_hint + "cond_buf", 16, 1, tir::BufferType::kDefault);
+    auto bool_cond = tir::Cast(DataType::Bool(1), tir::BufferLoad(cond_buffer, {tir::make_const(DataType::Int(32), 0)}));
+    auto if_else_stmts = temp_stmts_[GetRef<Expr>(op)];
+    auto if_stmt = tir::IfThenElse(bool_cond, tir::SeqStmt::Flatten(if_else_stmts.first), tir::SeqStmt::Flatten(if_else_stmts.second));
+    // stmts_.push_back(if_stmt);
+    // zz: need to handle nested if
+    if (!aot_codegen_if_expr_stack_.empty()) {
+      auto if_expr_scope = aot_codegen_if_expr_stack_.back();
+      if (if_expr_scope.second == 0) {
+        temp_stmts_[if_expr_scope.first].first.push_back(if_stmt);
+      }
+      else if (if_expr_scope.second == 1) {
+        temp_stmts_[if_expr_scope.first].second.push_back(if_stmt);
+      }
+      else {
+        ICHECK(false) << "if_expr_scope.second should be 0 or 1";
+      }
+    }
+    else {
+      stmts_.push_back(if_stmt);
+    }
   }
   void VisitExpr_(const FunctionNode* op) override {
     ICHECK(op->GetAttr<String>(attr::kCompiler).defined())
@@ -779,6 +1099,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   // runner function needs to be legalized by the LegalizePackedCalls pass.
   tir::PrimFunc CreateMainFunc(String mod_name, unsigned int relay_params) {
     tir::Stmt body = tir::SeqStmt::Flatten(stmts_);
+    // std::cout << "[ZZDEBUG] CreateMainFunc: Flatten body: " << std::endl << PrettyPrint(body) << std::endl;
     // Allocate the sids
     std::unordered_map<int, bool> allocated;
 
@@ -979,6 +1300,44 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     return lowered_mod;
   }
 
+  // zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+  IRModule PlanMemoryWithUSMPWithMainFuncName(const IRModule& mod, String main_func_name) {
+    VLOG(1) << "Planning memory with USMP for module:" << std::endl << PrettyPrint(mod);
+    Integer workspace_byte_alignment = GetModuleWorkspaceByteAlignment(mod);
+    IRModule lowered_mod = mod->ShallowCopy();
+    lowered_mod = tir::transform::UnifiedStaticMemoryPlannerWithMainFuncName(main_func_name)(lowered_mod); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main = "__tvm_main__"]
+    function_metadata_ = CalculateWorkspaceSizes(lowered_mod, function_metadata_);
+    Optional<Array<tir::usmp::AllocatedPoolInfo>> allocated_pool_infos =
+        lowered_mod->GetAttr<Array<tir::usmp::AllocatedPoolInfo>>(tvm::attr::kPoolArgs);
+    backend::FunctionInfo main_func_info =
+        lowered_mod->GetAttr<backend::FunctionInfo>("main_func_info").value();
+    main_func_info->workspace_sizes.clear();
+    if (allocated_pool_infos) {
+      for (const tir::usmp::AllocatedPoolInfo& allocated_pool_info : allocated_pool_infos.value()) {
+        for (const auto& tgt : allocated_pool_info->pool_info->targets) {
+          VLOG(1) << "USMP requires target " << tgt->ToDebugString() << " to have pool size "
+                  << allocated_pool_info->allocated_size->value;
+          size_t size = allocated_pool_info->allocated_size->value;
+          if (allocated_pool_info->pool_info->IsInstance<ConstantPoolInfoNode>()) {
+            size += main_func_info->constant_sizes.count(tgt)
+                        ? main_func_info->constant_sizes[tgt]->value
+                        : 0;
+            main_func_info->constant_sizes.Set(tgt, size);
+          } else if (allocated_pool_info->pool_info->IsInstance<WorkspacePoolInfoNode>()) {
+            size += main_func_info->workspace_sizes.count(tgt)
+                        ? main_func_info->workspace_sizes[tgt]->value
+                        : 0;
+            main_func_info->workspace_sizes.Set(tgt, size);
+          } else {
+            LOG(FATAL) << "Unknown pool type: " << allocated_pool_info->pool_info->GetTypeKey();
+          }
+        }
+      }
+    }
+    function_metadata_.Set(runtime::symbol::tvm_module_main, main_func_info);
+    return lowered_mod;
+  }
+
   /*!
    * \brief Run StorageRewrite to plan memory for lowered IRModule.
    */
@@ -1077,6 +1436,12 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   /*! \brief A set of variables that are let bound. */
   std::unordered_set<Var, ObjectPtrHash, ObjectPtrEqual> let_bound_vars_;
 
+  // zz: for obfuscation function name
+  int obf_func_cnt_ = 0;
+  // zz: for if handling
+  std::vector<std::pair<Expr, int> > aot_codegen_if_expr_stack_; // if_expr, current scope (0 for if_branch, 1 for else_branch)
+  std::unordered_map<Expr, std::pair<std::vector<tir::Stmt>, std::vector<tir::Stmt> >, runtime::ObjectPtrHash, runtime::ObjectPtrEqual> temp_stmts_; // if_expr -> (if_branch_stmts, else_branch_stmts)
+
  public:
   AOTExecutorCodegen(runtime::Module* mod, const Array<Target>& targets)
       : mod_(mod), config_(transform::PassContext::Current(), targets) {}
@@ -1121,12 +1486,96 @@ class AOTExecutorCodegen : public MixedModeVisitor {
                     << ") is not one of the expected values";
     }
 
+    // std::cout << "[AOTExecutorCodegen::Codegen] before ToANormalForm mod:" << std::endl << PrettyPrint(mod) << "[AOTExecutorCodegen::Codegen] before ToANormalForm mod end." << std::endl;
     mod = transform::ToANormalForm()(mod);
+    // std::cout << "[AOTExecutorCodegen::Codegen] after ToANormalForm mod:" << std::endl << PrettyPrint(mod) << "[AOTExecutorCodegen::Codegen] after ToANormalForm mod end." << std::endl;
     mod = transform::InferType()(mod);
-    mod = transform::AnnotateUsedMemory()(mod);
+    // std::cout << "[AOTExecutorCodegen::Codegen] after InferType mod:" << std::endl << PrettyPrint(mod) << "[AOTExecutorCodegen::Codegen] after InferType mod end." << std::endl;
+    // mod = transform::AnnotateUsedMemory()(mod); // zz: seems to be unuseful
+
+    /* =========================================================================================================
+       * [zz] main logic:
+       * perfrom obf funtion annotation
+       * call CodegenObfCallSelected to generate obfuscated function
+       ========================================================================================================= */
+    // ----------------- zz: collect non-prim call function
+    auto main_func = Downcast<Function>(mod->Lookup("main"));
+    NonPrimCallCollector non_prim_call_collector;
+    non_prim_call_collector.VisitExpr(main_func->body);
+    std::vector<Function> original_non_prim_funcs = non_prim_call_collector.GetNonPrimCallFuncs();
+
+    // ----------------- zz: do non-prim call annotation
+    NonPrimCallAnnotator non_prim_call_annotator(original_non_prim_funcs);
+    auto update_main_func = non_prim_call_annotator.Mutate(main_func);
+    // std::vector<Function> annotated_non_prim_funcs = non_prim_call_annotator.GetNonPrimCallFuncs();
+    std::vector<std::pair<int, Function> > non_prim_obfid_funcs = non_prim_call_annotator.GetNonPrimCallFuncs();
+    mod->Add(mod->GetGlobalVar("main"), Downcast<Function>(update_main_func));
+
+    std::unordered_map<int, SubModItem> func_to_replace;
+
+    // ----------------- zz: lower obfuscated function using CodegenObfCallSelected
+    // for (auto func : annotated_non_prim_funcs) {
+    // for (size_t func_idx = 0; func_idx < original_non_prim_funcs.size(); func_idx++) {
+    for (auto obfid_func : non_prim_obfid_funcs) {
+      int obf_func_idx = obfid_func.first;
+      Function func_orig = obfid_func.second;
+      // auto func_orig = original_non_prim_funcs[func_idx];
+      // auto func_annot = annotated_non_prim_funcs[func_idx];
+      std::string obf_func_mod_name = "obf" + std::to_string(obf_func_idx);
+      String obf_func_name = String(obf_func_mod_name);
+      auto func_for_mod = WithAttrs(func_orig, {{tvm::attr::kGlobalSymbol, obf_func_name}});
+      auto non_prim_mod = IRModule::FromExpr(func_for_mod);
+      non_prim_mod = WithAttrs(non_prim_mod, {{tvm::attr::kRuntime, mod->GetAttr<Runtime>(tvm::attr::kRuntime).value()}, {tvm::attr::kExecutor, mod->GetAttr<Executor>(tvm::attr::kExecutor).value()}});
+      String obf_mod_name = String(obf_func_mod_name);
+      auto non_prim_mod_ret = CodegenObfCallSelected(non_prim_mod, func_for_mod, obf_mod_name);
+      // auto non_prim_mod_ret = CodegenObfCall(non_prim_mod, func_for_mod, obf_mod_name);
+      CHECK_EQ(non_prim_mod_ret.lowered_funcs.size(), 1);
+      for (auto target_mod : non_prim_mod_ret.lowered_funcs) {
+        auto obf_mod = target_mod.second;
+        // func_to_replace.insert(std::make_pair(func.get(), obf_mod.get()));
+        func_to_replace[obf_func_idx] = std::make_pair(obf_func_name, obf_mod);
+        // for (auto obf_mod_func : obf_mod->functions) {
+        //   auto global_var_supply = GlobalVarSupply(mod);
+        //   mod->Add(global_var_supply->FreshGlobal(obf_mod_func->, false), obf_mod_func);
+        // }
+        // for (auto omgv : obf_mod->GetGlobalVars()) {
+        //   // mod->Add(omgv, obf_mod->Lookup(omgv));
+        // }
+        // for (const auto& kv : obf_mod->functions) {
+        //   const GlobalVar& var = kv.first;
+        //   const BaseFunc& func = kv.second;
+        //   mod->Update(var, func);
+        // }
+      }
 
+      // ----------------- zz: clear structures used in CodegenObfCallSelected
+      // TODO
+      storage_device_map_.clear();
+      input_vars_.clear();
+      return_sid_.clear();
+      // the main function signatures
+      main_signature_.clear();
+      main_buffer_map_.clear();
+      io_tensor_types_.clear();
+      // TODO
+      sids_table_.clear();
+      let_bound_vars_.clear();
+      constant_map_.clear(); // [TODO] zz: check if this is needed
+      params_.clear(); // [TODO] zz: params_ seems to be empty in example
+
+      // generated main function stmts
+      stmts_.clear();
+
+      // additional structures for obfuscation
+      aot_codegen_if_expr_stack_.clear();
+      temp_stmts_.clear();
+    }
+    /* =========================================================================================================
+       * [zz] main logic end
+       ========================================================================================================= */
     IRModule lowered_mod =
-        tec::LowerTE(mod_name, config_, [this, workspace_byte_alignment](BaseFunc func) {
+        // tec::LowerTE(mod_name, config_, [this, workspace_byte_alignment](BaseFunc func) {
+        tec::LowerTE_Simple(mod_name, config_, [this, workspace_byte_alignment](BaseFunc func) {
           // We need to maintain the constant map for external
           // functions so we pass this processing function which
           // allows us to process each function as we lower it.
@@ -1138,14 +1587,21 @@ class AOTExecutorCodegen : public MixedModeVisitor {
           // execute as a further pass, instead writing data to the
           // lowering process directly.
           tec::UpdateFunctionMetadata(func, this->function_metadata_, workspace_byte_alignment);
-        })(mod);
+        }, func_to_replace)(mod);
 
     transform::PassContext pass_ctx = transform::PassContext::Current();
-    bool enable_remove_reshapes =
-        pass_ctx->GetConfig<Bool>("relay.remove_standalone_reshapes.enable", Bool(true)).value();
-    if (enable_remove_reshapes) {
-      lowered_mod = transform::RemoveStandaloneReshapes()(lowered_mod);
-    }
+    // -------------------------------------------------------------------------------------------------
+    // comment out: zz: this pass causes removal of prim function called by non-prim (obf) functions
+    // -------------------------------------------------------------------------------------------------
+    // bool enable_remove_reshapes =
+    //     pass_ctx->GetConfig<Bool>("relay.remove_standalone_reshapes.enable", Bool(true)).value();
+    // if (enable_remove_reshapes) {
+    //   lowered_mod = transform::RemoveStandaloneReshapes()(lowered_mod);
+    //   std::cout << "[AOTExecutorCodegen::Codegen] lowered_mod after RemoveStandaloneReshapes mod:" << std::endl << PrettyPrint(lowered_mod) << std::endl << "[AOTExecutorCodegen::Codegen] lowered_mod after RemoveStandaloneReshapes mod end." << std::endl;
+    // }
+    // -------------------------------------------------------------------------------------------------
+    // comment out end
+    // -------------------------------------------------------------------------------------------------
     auto lowered_main = lowered_mod->Lookup("main");
     auto lowered_main_func = Downcast<Function>(lowered_main);
 
@@ -1348,6 +1804,299 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     return ret;
   }
 
+  // zz: for non-prim obf function codegen
+  LoweredOutput CodegenObfCallSelected(IRModule mod, relay::Function func, String mod_name) {
+  //   VLOG_CONTEXT << "AOT";
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected input mod: " << std::endl << PrettyPrint(mod) << std::endl << "[ZZDEBUG] CodegenObfCallSelected input mod end." << std::endl;
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected input func: " << std::endl << PrettyPrint(func) << std::endl << "[ZZDEBUG] CodegenObfCallSelected input func end." << std::endl;
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected input mod_name: " << mod_name << std::endl;
+
+    Runtime runtime_config = mod->GetAttr<Runtime>(tvm::attr::kRuntime).value();
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected runtime_config end." << std::endl;
+
+    Integer workspace_byte_alignment = GetModuleWorkspaceByteAlignment(mod);
+
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected workspace_byte_alignment end." << std::endl;
+
+    Executor executor_config = mod->GetAttr<Executor>(tvm::attr::kExecutor).value();
+    std::string interface_api =
+        executor_config->GetAttr<String>("interface-api").value_or("packed");
+    bool unpacked_api = executor_config->GetAttr<Bool>("unpacked-api").value_or(Bool(false));
+
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected unpacked_api end." << std::endl;
+
+    // Validate choice of unpacked_api and use_call_cpacked_
+    if (runtime_config->name == kTvmRuntimeCrt) {
+      if (unpacked_api == true) {
+        call_type_ = CallType::kUnpacked;
+      } else if (unpacked_api == false && interface_api == "packed") {
+        call_type_ = CallType::kCPacked;
+      } else {
+        CHECK(interface_api == "packed" || unpacked_api == true)
+            << "Either need interface_api == \"packed\" (got: " << interface_api
+            << ") or unpacked-api == true (got: " << unpacked_api << ") when targeting c runtime";
+        ICHECK(false) << "Unhandled executor option config: interface-api=" << interface_api
+                      << ", unpacked-api=" << unpacked_api;
+      }
+    } else if (runtime_config->name == kTvmRuntimeCpp) {
+      if (unpacked_api == false && interface_api == "packed") {
+        call_type_ = CallType::kCPacked;
+      } else {
+        CHECK(static_cast<bool>(unpacked_api) == false && interface_api == "packed")
+            << "Need unpacked-api == false (got: " << unpacked_api
+            << ") and interface-api == \"packed\" (got: " << interface_api
+            << ") when targeting c++ runtime";
+        ICHECK(false) << "Unhandled executor option config: interface-api=" << interface_api
+                      << ", unpacked-api=" << unpacked_api;
+      }
+    } else {
+      ICHECK(false) << "runtime_config (" << runtime_config->name
+                    << ") is not one of the expected values";
+    }
+
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected Validate end." << std::endl;
+
+    // -----------------------------------------------------------------
+    // zz: ToANormalForm and InferType are done in Codegen
+    // -----------------------------------------------------------------
+    // mod = transform::ToANormalForm()(mod);
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after ToANormalForm mod:" << std::endl << PrettyPrint(mod) << "[ZZDEBUG] CodegenObfCallSelected after ToANormalForm mod end." << std::endl;
+    // mod = transform::InferType()(mod);
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after InferType mod:" << std::endl << PrettyPrint(mod) << "[ZZDEBUG] CodegenObfCallSelected after InferType mod end." << std::endl;
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected before AnnotateUsedMemoryWithMainFuncName mod:" << std::endl << PrettyPrint(mod) << "[ZZDEBUG] CodegenObfCallSelected before AnnotateUsedMemoryWithMainFuncName mod end." << std::endl;
+    mod = transform::AnnotateUsedMemoryWithMainFuncName(mod_name)(mod); // zz: [originally operate on "main"]
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after AnnotateUsedMemoryWithMainFuncName mod:" << std::endl << PrettyPrint(mod) << "[ZZDEBUG] CodegenObfCallSelected after AnnotateUsedMemoryWithMainFuncName mod end." << std::endl;
+
+    IRModule lowered_mod =
+        // tec::LowerTE(mod_name, config_, [this, workspace_byte_alignment](BaseFunc func) {
+        tec::LowerTE_Simple(mod_name, config_, [this, workspace_byte_alignment](BaseFunc func) {
+          // We need to maintain the constant map for external
+          // functions so we pass this processing function which
+          // allows us to process each function as we lower it.
+          if (func->GetAttr<String>(attr::kCompiler).defined()) {
+            UpdateConstants(func, &params_);
+          }
+
+          // TODO(@areusch, @jroesch): We should refactor this to
+          // execute as a further pass, instead writing data to the
+          // lowering process directly.
+          tec::UpdateFunctionMetadata(func, this->function_metadata_, workspace_byte_alignment);
+        })(mod);
+
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after LowerTE mod:" << std::endl << PrettyPrint(lowered_mod) << "[ZZDEBUG] CodegenObfCallSelected after LowerTE mod end." << std::endl;
+
+    transform::PassContext pass_ctx = transform::PassContext::Current();
+    // zz: the RemoveStandaloneReshapesWithMainFuncName may have influence
+    // bool enable_remove_reshapes =
+    //     pass_ctx->GetConfig<Bool>("relay.remove_standalone_reshapes.enable", Bool(true)).value();
+    // if (enable_remove_reshapes) {
+    //   std::cout << "[ZZDEBUG] CodegenObfCallSelected enable_remove_reshapes." << std::endl;
+    //   lowered_mod = transform::RemoveStandaloneReshapesWithMainFuncName(mod_name)(lowered_mod); // zz: [originally operate on "main"]
+    // }
+    auto lowered_main = lowered_mod->Lookup(mod_name); // zz: [originally "main"]
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected lowered_main" << std::endl << PrettyPrint(lowered_main) << std::endl;
+
+    auto lowered_main_func = Downcast<Function>(lowered_main);
+
+    // zz: get let_bindings for AOTOnDemandAllocator
+    // std::cout << "[AOTExecutorCodegen::Codegen] lowered_main_func:" << std::endl << PrettyPrint(lowered_main_func) << std::endl;
+    // LetBindingCollector let_binding_collector;
+    // let_binding_collector.VisitExpr(lowered_main_func->body);
+    // std::unordered_map<Expr, Expr, runtime::ObjectPtrHash, runtime::ObjectPtrEqual> let_bound_vars = let_binding_collector.GetLetBindings();
+
+    // Post-lowering storage map for writing main func
+    AOTOnDemandAllocator final_aot_allocator;
+    final_aot_allocator.Run(lowered_main_func);
+    storage_device_map_ = final_aot_allocator.GetStorageMap();
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after AOTOnDemandAllocator" << std::endl;
+    // for (auto kv : storage_device_map_) {
+    //   std::cout << "[ZZDEBUG] CodegenObfCallSelected storage_device_map_ key expr:" << std::endl << PrettyPrint(kv.first) << std::endl;
+    //   std::cout << "[ZZDEBUG] CodegenObfCallSelected storage_device_map_ storage id:" << std::endl << PrettyPrint(kv.second) << std::endl << std::endl;
+    //   // for (auto sid : kv.second->storage_ids) {
+    //   //   std::cout << "  (zz) storage_id:" << sid << std::endl;
+    //   // }
+    // }
+    // final_aot_allocator.PrintStorageDeviceMap();
+
+    // TODO(@electriclilies, @jroesch, @Mousius): remove UpdateMainWorkspaceSize
+    StaticMemoryPlan memory_plan(storage_device_map_);
+    backend::FunctionInfo func_info =
+        tec::UpdateMainWorkspaceSize(lowered_mod, config_, memory_plan->expr_to_storage_info, mod_name); // zz: [originally operate on "main"]
+    lowered_mod = WithAttr(lowered_mod, "main_func_info", func_info);
+    // std::cout << "[AOTExecutorCodegen::Codegen] lowered_mod after UpdateMainWorkspaceSize mod:" << std::endl << PrettyPrint(lowered_mod) << "[AOTExecutorCodegen::Codegen] lowered_mod after UpdateMainWorkspaceSize mod end." << std::endl;
+
+
+    for (auto input : lowered_main_func->params) {
+      input_vars_.push_back(input);
+      std::string input_name = SanitizeName(input->name_hint());
+      // We dont want the compiler changing input names in the
+      // event of a sanitization collision. Therefore, enforcing
+      // the var created to use the input_name strictly.
+      CreateIOVar(input, input_name, /*use_unique_name = */ false);
+    }
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after SanitizeName" << std::endl;
+
+    // Define the storage allocator ids
+    for (auto kv : storage_device_map_) {
+      for (auto sid : kv.second->storage_ids) {
+        // The buffer_var is created with storage_scope to be global.workspace to be serviced by
+        // TVMBackendAllocWorkspace(TVMBAW) calls, explicitly. The reasoning being the executor
+        // allocates should be serviced by TVMBAWs as the data could be accessed by many devices and
+        // should not be lowered to the stack. For more details please refer to the discussion here:
+        // https://github.com/apache/tvm/issues/9022
+        te::Var buffer_var(MakeString("sid_", sid),
+                           PointerType(PrimType(DataType::Int(8)), "global.workspace"));
+        sids_table_[sid] = buffer_var;
+      }
+    }
+
+    // Retrieve the return sids
+    return_sid_ = final_aot_allocator.GetReturnIds();
+    // Insert outputs to main func signature
+    // If output tensor names were provided use them
+    if (auto opt = func->GetAttr<Array<String>>("output_tensor_names")) {
+      Array<String> output_tensor_names = opt.value();
+      Expr output_expr = lowered_main_func->body;
+      if (output_expr->checked_type()->IsInstance<TupleTypeNode>()) {
+        TupleType output_tuple_type = Downcast<TupleType>(output_expr->checked_type());
+        for (unsigned i = 0; i < output_tuple_type->fields.size(); i++) {
+          // AoT Executor Codegen does not create these names,
+          // thus should be used as they are provided.
+          CreateIOVar(output_tuple_type->fields[i], output_tensor_names[i],
+                      /*use_unique_name = */ false);
+        }
+      } else {
+        // AoT Executor Codegen does not create these names,
+        // thus should be used as they are provided.
+        CreateIOVar(lowered_main_func->body, output_tensor_names[0], /*use_unique_name = */ false);
+      }
+    } else {
+      // If output tensor names are not provided we will generate output(x)
+      // where x is a counter to create unique names.
+      CreateIOVar(lowered_main_func->body, "output");
+    }
+
+    CollectDeviceVariables(lowered_mod->GetAttr<Map<GlobalVar, String>>("device_contexts").value());
+    num_arguments_ = [&]() -> Map<GlobalVar, Integer> {
+      Map<GlobalVar, Integer> arg_count;
+      for (const auto& [gvar, func] : lowered_mod->functions) {
+        if (const auto* prim_func = func.as<tir::PrimFuncNode>()) {
+          arg_count.Set(gvar, prim_func->params.size());
+        }
+      }
+      return arg_count;
+    }();
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected before VisitExpr lowered_main_func:" << std::endl << PrettyPrint(lowered_main_func) << std::endl << "[ZZDEBUG] CodegenObfCallSelected before VisitExpr lowered_main_func end." << std::endl;
+    VisitExpr(lowered_main_func->body);
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after VisitExpr lowered_main_func:" << std::endl << PrettyPrint(lowered_main_func) << std::endl << "[ZZDEBUG] CodegenObfCallSelected after VisitExpr lowered_main_func end." << std::endl;
+
+    // Create the runner function. Please note that the function is not legal yet
+    // because the packed calls arguments are not wrapped in TVMValues. To make this happen we need
+    // to run the LegalizePackedCalls pass.
+    LoweredOutput ret;
+
+    // Collect any constants extracted by external codegen.
+    ret.params = std::unordered_map<std::string, tvm::runtime::NDArray>();
+    Map<String, runtime::NDArray> const_name_to_constant =
+        lowered_mod->GetAttr<Map<String, runtime::NDArray>>(tvm::attr::kConstNameToConstant)
+            .value_or({});
+    for (const auto& kv : const_name_to_constant) {
+      ICHECK(ret.params.emplace(kv.first, kv.second).second);
+    }
+
+    // Collect any constants extracted during lowering.
+    for (const auto& kv : params_) {
+      ICHECK(ret.params.emplace(kv.first, kv.second).second);
+    }
+
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after constants collection" << std::endl;
+
+    // AoT Executor codegen works completely on TIR beyond this point, hence removing relay main
+    // function and replacing it with its TIR version. We should try to make this a Pass.
+    auto lowered_mod_main_globvar = lowered_mod->GetGlobalVar(mod_name);
+    auto lowered_mod_main_globvar_func = lowered_mod->Lookup(lowered_mod_main_globvar);
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected lowered_mod before Remove lowered_mod_main_globvar type: " << PrettyPrint(lowered_mod_main_globvar->checked_type_) << std::endl;
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected lowered_mod before Remove lowered_mod_main_globvar_func type: " << PrettyPrint(lowered_mod_main_globvar_func->checked_type_) << std::endl;
+    auto stored_type = lowered_mod_main_globvar_func->checked_type_;
+    lowered_mod->Remove(lowered_mod->GetGlobalVar(mod_name)); // zz: [originally "main"]
+    auto tir_main_func = CreateMainFunc(mod_name, lowered_main_func->params.size());
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected CreateMainFunc:" << std::endl << PrettyPrint(tir_main_func) << "[ZZDEBUG] CodegenObfCallSelected CreateMainFunc end." << std::endl;
+    // Extract additional information around main TIR PrimFunc arguments
+    Array<String> devices = ListDevices();
+    const auto main_func_params_end_iterator =
+        tir_main_func->params.begin() + tir_main_func->params.size();
+    const auto outputs_begin_iterator =
+        main_func_params_end_iterator - return_sid_.size() - devices.size();
+    Array<tir::Var> inputs = Array<tir::Var>(tir_main_func->params.begin(), outputs_begin_iterator);
+    Array<TensorType> input_tensor_types;
+    for (auto i : inputs) {
+      input_tensor_types.push_back(io_tensor_types_[i]);
+    }
+    Array<tir::Var> outputs =
+        Array<tir::Var>(outputs_begin_iterator, main_func_params_end_iterator - devices.size());
+
+    tir_main_func = WithAttr(std::move(tir_main_func), "global_symbol", runtime::String(mod_name)); // zz: the name is profixed with ___tvm_main__
+    lowered_mod->Update(GlobalVar(mod_name), tir_main_func); // zz: [originally ::tvm::runtime::symbol::tvm_module_main = "__tvm_main__"]
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after Update tir_main_func:" << std::endl << PrettyPrint(lowered_mod) << std::endl << "[ZZDEBUG] CodegenObfCallSelected after Update tir_main_func end." << std::endl;
+    auto tir_main_globvar = lowered_mod->GetGlobalVar(mod_name);
+    auto tir_main_globvar_func = lowered_mod->Lookup(tir_main_globvar);
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected lowered_mod after Update tir_main_globvar type: " << PrettyPrint(tir_main_globvar->checked_type_) << std::endl;
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected lowered_mod after Update tir_main_globvar_func type: " << PrettyPrint(tir_main_globvar_func->checked_type_) << std::endl;
+    tir_main_globvar_func->checked_type_ = stored_type;
+
+    // Parallel for loops are not supported in AoT codegen.
+    lowered_mod = tir::transform::ConvertForLoopsToSerial()(lowered_mod);
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected after ConvertForLoopsToSerial" << std::endl;
+
+    // // Check USMP option
+    // bool enable_usmp = false;
+    // if (runtime_config->name == kTvmRuntimeCrt) {
+    //   enable_usmp = true;
+    // }
+    // if (pass_ctx->GetConfig<Bool>(kUSMPEnableOption) != nullptr) {
+    //   enable_usmp = pass_ctx->GetConfig<Bool>(kUSMPEnableOption, Bool(false)).value();
+    // }
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected enable_usmp=" << enable_usmp << std::endl;
+
+    // if (enable_usmp) {
+    //   std::cout << "[ZZDEBUG] CodegenObfCallSelected before PlanMemoryWithUSMP" << std::endl;
+    //   lowered_mod = PlanMemoryWithUSMPWithMainFuncName(lowered_mod, mod_name); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main = "__tvm_main__" in UnifiedStaticMemoryPlanner]
+    //   std::cout << "[ZZDEBUG] CodegenObfCallSelected after PlanMemoryWithUSMP" << std::endl;
+    // } else {
+    //   lowered_mod = PlanMemoryWithStorageRewrite(lowered_mod);
+    //   std::cout << "[ZZDEBUG] CodegenObfCallSelected after PlanMemoryWithStorageRewrite" << std::endl;
+    // }
+    // ret.function_metadata = std::move(function_metadata_);
+
+    // // Legalize AOT if needed. This means that all the packed calls
+    // // need to be wrapped in TVMValues (unless unpacked_api is set)
+    // if (call_type_ == CallType::kCPacked || call_type_ == CallType::kPacked) {
+    //   auto pack_calls = tir::transform::LegalizePackedCalls();
+    //   lowered_mod = pack_calls(lowered_mod);
+    // }
+
+    // // Collect any runtime modules generated by external codegen.
+    // ret.external_mods =
+    //     lowered_mod->GetAttr<Array<tvm::runtime::Module>>(tvm::attr::kExternalMods).value_or({});
+
+    // This is the point where we separate the functions in the module by target
+    VLOG(1) << "lowered module:" << std::endl << PrettyPrint(lowered_mod);
+    ret.lowered_funcs = tec::GetPerTargetModules(lowered_mod);
+    tir_main_globvar = lowered_mod->GetGlobalVar(mod_name);
+    tir_main_globvar_func = lowered_mod->Lookup(tir_main_globvar);
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected lowered_mod after ret.lowered_funcs tir_main_globvar type: " << PrettyPrint(tir_main_globvar->checked_type_) << std::endl;
+    // std::cout << "[ZZDEBUG] CodegenObfCallSelected lowered_mod after ret.lowered_funcs tir_main_globvar_func type: " << PrettyPrint(tir_main_globvar_func->checked_type_) << std::endl;
+    VLOG(1) << "per-target modules:";
+    for (const auto& kv : ret.lowered_funcs) {
+      VLOG(1) << "target:" << std::endl
+              << kv.first->ToDebugString() << std::endl
+              << "maps to:" << std::endl
+              << PrettyPrint(kv.second);
+    }
+
+    return ret;
+  }
+
   /*!
    * \brief Get list of devices found
    * \return List of devices
diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc
index 83c252d83..39af761df 100644
--- a/src/relay/backend/build_module.cc
+++ b/src/relay/backend/build_module.cc
@@ -43,12 +43,41 @@
 #include "utils.h"
 
 namespace tvm {
+// zz: config option for dnnobf
+// TVM_REGISTER_PASS_CONFIG_OPTION("relay.dnnobf", Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION("dnnobf.profile", Bool); // zz: dnnobf profile option
+TVM_REGISTER_PASS_CONFIG_OPTION("dnnobf.flex_fuse", Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION("dnnobf.fakeop_insert", Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION("dnnobf.compute_reorder", Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION("dnnobf.selected_neuron_indices", Array<Array<Integer>>); // zz: selected neuron indices for obfuscated ops
+TVM_REGISTER_PASS_CONFIG_OPTION("dnnobf.selected_neuron_ranges", Array<Array<FloatImm>>); // zz: neuron ranges for obfuscated ops
 namespace relay {
 namespace transform {
 Pass LabelOps();
 }
 namespace backend {
 
+// zz: visitor for extract non prim obfuscated calls in main function
+class ObfAttrTester : public ExprVisitor {
+ public:
+  ObfAttrTester() {}
+
+  void VisitExpr_(const CallNode* call_node) final {
+    ExprVisitor::VisitExpr(call_node->op);
+    for (auto arg : call_node->args) {
+      ExprVisitor::VisitExpr(arg);
+    }
+    if (call_node->obf_attrs.defined() && call_node->obf_attrs.HasIntAttr("obf_op_id")) {
+      int obf_op_id = call_node->obf_attrs.GetAttr<Integer>("obf_op_id").value().IntValue();
+      std::cout << "[ZZDEBUG] ObfAttrTester::VisitExpr_ obf_op_id: " << obf_op_id << std::endl;
+      Array<FloatImm> neuron_ranges = call_node->obf_attrs.GetAttr<Array<FloatImm>>("neuron_ranges").value();
+      for (auto neuron_range : neuron_ranges) {
+        std::cout << "    (zz) ObfAttrTester::VisitExpr_ neuron_range: " << neuron_range->value << std::endl;
+      }
+    }
+  }
+};
+
 using namespace tvm::relay::transform;
 
 /*!
@@ -327,13 +356,35 @@ class RelayBuildModule : public runtime::ModuleNode {
 
   IRModule OptimizeImpl(IRModule relay_module) {
     ICHECK(relay_module.defined()) << "The IRModule must be defined for the Relay compiler.";
+    // zz: dnnobf configs
+    Function func_test_obf_attr = Downcast<Function>(relay_module->Lookup("main"));
+    transform::PassContext pass_ctx = transform::PassContext::Current();
+    // bool enable_dnnobf = pass_ctx->GetConfig<Bool>("relay.dnnobf", Bool(false)).value();
+    bool is_profile_build = pass_ctx->GetConfig<Bool>("dnnobf.profile", Bool(false)).value();
+    bool enable_flex_fuse = pass_ctx->GetConfig<Bool>("dnnobf.flex_fuse", Bool(false)).value();
+    bool enable_fakeop_insert = pass_ctx->GetConfig<Bool>("dnnobf.fakeop_insert", Bool(false)).value();
+    if (is_profile_build) {
+      if (enable_flex_fuse || enable_fakeop_insert) {
+        LOG(FATAL) << "dnnobf.flex_fuse or dnnobf.fakeop_insert should not be set when dnnobf.profile is true";
+      }
+    }
+    if (enable_fakeop_insert) {
+      ICHECK(enable_flex_fuse) << "dnnobf.flex_fuse need to be set if dnnobf.fakeop_insert is set";
+      std::cout << "[ZZDEBUG] BuildRelay enable_fakeop_insert = true" << std::endl;
+      ObfAttrTester obf_attr_tester;
+      obf_attr_tester.VisitExpr(func_test_obf_attr->body);
+    }
 
     backend::BindParamsInModule(relay_module, params_);
 
     Array<Pass> pass_seqs =
         GetPassPrefix(/*is_homogenous=*/config_->primitive_targets.size() == 1, /*is_vm=*/false);
-    transform::PassContext pass_ctx = PassContext::Current();
+    // transform::PassContext pass_ctx = PassContext::Current();
+    std::cout << "[INFO] OptimizeImpl pass_seqs.size() GetPassPrefix = " << pass_seqs.size() << std::endl;
 
+    // ====================================================
+    // zz: comment out this pass not influence the example
+    // ====================================================
     if (config_->optional_homogeneous_target.defined()) {
       // This pass currently only supports the homogeneous case.
       pass_seqs.push_back(transform::SplitArgs(
@@ -341,25 +392,59 @@ class RelayBuildModule : public runtime::ModuleNode {
               .value()
               .IntValue()));
     }
+    std::cout << "[INFO] OptimizeImpl pass_seqs.size() optional_homogeneous_target = " << pass_seqs.size() << std::endl;
+    // ==========================================================
+    // zz: end - comment out pass part not influence the example
+    // ==========================================================
 
     // Always plan devices so the remaining passes don't need to distinguish homogeneous vs
     // hetrogenous execution.
     pass_seqs.push_back(transform::PlanDevices(config_));
+    std::cout << "[INFO] OptimizeImpl pass_seqs.size() PlanDevices = " << pass_seqs.size() << std::endl;
 
     // Fuse the operations if it is needed.
-    pass_seqs.push_back(transform::FuseOps());
+    if (is_profile_build) {
+      pass_seqs.push_back(transform::InterestingExprExaminer());
+      pass_seqs.push_back(transform::PlanDevices(config_)); // zz: need to plan devices again
+    }
+    if (!enable_flex_fuse && !enable_fakeop_insert) {
+      pass_seqs.push_back(transform::FuseOps());
+    }
+    // pass_seqs.push_back(transform::FuseOps());
+    // pass_seqs.push_back(transform::FuseOps(0)); // zz: set opt_level to 0, avoid fuse
+    else {
+      pass_seqs.push_back(transform::AnnotateObfAttrs());
+      pass_seqs.push_back(transform::FuseOpsObf(-1, enable_flex_fuse, enable_fakeop_insert));
+      pass_seqs.push_back(transform::PlanDevices(config_)); // zz: need to plan devices again for inserted fake ops
+      // set realy.dnnobf to false for the following passes
+      // add infertype pass
+      pass_seqs.push_back(transform::InferType());
+      if (enable_fakeop_insert) {
+        pass_seqs.push_back(transform::FuseOps(0)); // zz: need to add another fuse for non-primitive subfunctions
+      }
+    }
+    std::cout << "[INFO] OptimizeImpl pass_seqs.size() FuseOps = " << pass_seqs.size() << std::endl;
 
     // Create a sequential pass and perform optimizations.
-    transform::Pass seq = transform::Sequential(pass_seqs);
+    // transform::Pass seq = transform::Sequential(pass_seqs);
+    transform::Pass seq = transform::Sequential(pass_seqs, "ZZDEBUG_OptimizeImpl");
     if (config_->optional_homogeneous_target.defined()) {
+      // std::cout << "----- OptimizeImpl before Sequential mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl before Sequential mod end" << std::endl;
+      // std::cout << "[Option] OptimizeImpl optional_homogeneous_target." << std::endl;
       With<Target> tctx(config_->optional_homogeneous_target);
       relay_module = seq(relay_module);
+      // std::cout << "----- OptimizeImpl after Sequential mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl after Sequential mod end" << std::endl;
     } else {
+      // std::cout << "[Option] OptimizeImpl NOT optional_homogeneous_target (not in exmaple model)." << std::endl;
       relay_module = seq(relay_module);
     }
 
+    // ====================================================
+    // zz: comment out this pass not influence the example
+    // ====================================================
     // Do layout rewrite for auto-scheduler.
     if (backend::IsAutoSchedulerEnabled() && config_->optional_homogeneous_target.defined()) {
+      std::cout << std::endl << "[ZZDEBUG] OptimizeImpl IsAutoSchedulerEnabled (not in exmaple model)"  << std::endl;
       Pass major_pass = transform::AutoSchedulerLayoutRewrite();
       bool enable_layout_rewrite_targets =
           config_->optional_homogeneous_target->GetTargetDeviceType() == kDLCPU ||
@@ -374,6 +459,7 @@ class RelayBuildModule : public runtime::ModuleNode {
       }
     }
     if (backend::IsMetaScheduleEnabled() && config_->optional_homogeneous_target.defined()) {
+      std::cout << std::endl << "[ZZDEBUG] OptimizeImpl IsMetaScheduleEnabled (not for lenet)"  << std::endl;
       Pass major_pass = transform::MetaScheduleLayoutRewrite();
       bool enable_layout_rewrite_targets =
           config_->optional_homogeneous_target->GetTargetDeviceType() == kDLCPU ||
@@ -387,8 +473,13 @@ class RelayBuildModule : public runtime::ModuleNode {
         relay_module = transform::FuseOps()(relay_module);
       }
     }
+    // ==========================================================
+    // zz: end - comment out pass part not influence the example
+    // ==========================================================
 
+    // std::cout << "----- OptimizeImpl before InferType mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl before InferType mod end" << std::endl;
     relay_module = transform::InferType()(relay_module);
+    // std::cout << "----- OptimizeImpl after InferType mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl after InferType mod end" << std::endl;
 
     // Inline the functions that have been lifted by the module scope.
     //
@@ -397,12 +488,15 @@ class RelayBuildModule : public runtime::ModuleNode {
     // inline functions. However, this should be very unlikely for accelerators
     // and vendor-provided libraries. So we don't handle for now.
     relay_module = transform::Inline()(relay_module);
+    // std::cout << "----- OptimizeImpl after Inline mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl after Inline mod end" << std::endl;
     relay_module = transform::InferType()(relay_module);
+    // std::cout << "----- OptimizeImpl after InferType mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl after InferType mod end" << std::endl;
     relay_module = transform::LabelOps()(relay_module);
+    // std::cout << "----- OptimizeImpl after LabelOps mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl after LabelOps mod end" << std::endl;
     relay_module = transform::AnnotateMemoryScope()(relay_module);
+    // std::cout << "----- OptimizeImpl after AnnotateMemoryScope mod:" << std::endl << PrettyPrint(relay_module) << "----- OptimizeImpl after AnnotateMemoryScope mod end" << std::endl;
 
     ICHECK(relay_module.defined());
-
     return relay_module;
   }
 
@@ -413,10 +507,25 @@ class RelayBuildModule : public runtime::ModuleNode {
    * \param params The parameters.
    */
   void BuildRelay(IRModule relay_module, const String& mod_name) {
+    // // zz: test dnnobf
+    // Function func_test_obf_attr = Downcast<Function>(relay_module->Lookup("main"));
+    // transform::PassContext pass_ctx = transform::PassContext::Current();
+    // bool enable_dnnobf = pass_ctx->GetConfig<Bool>("relay.dnnobf", Bool(false)).value();
+    // if (enable_dnnobf) {
+    //   std::cout << "[ZZDEBUG] BuildRelay enable_dnnobf = true" << std::endl;
+    //   ObfAttrTester obf_attr_tester;
+    //   obf_attr_tester.VisitExpr(func_test_obf_attr->body);
+    // }
+
     // Relay IRModule -> IRModule optimizations.
+    std::cout << "BuildRelay input mod:" << std::endl << PrettyPrint(relay_module) << "BuildRelay input mod end." << std::endl;
+    std::cout << "================= OptimizeImpl =================" << std::endl;
     IRModule module = WithAttrs(
         relay_module, {{tvm::attr::kExecutor, executor_}, {tvm::attr::kRuntime, runtime_}});
     relay_module = OptimizeImpl(std::move(module));
+    // relay_module = std::move(module);
+    std::cout << "=============== OptimizeImpl end ===============" << std::endl;
+    std::cout << "BuildRelay after OptimizeImpl mod:" << std::endl << PrettyPrint(relay_module) << "BuildRelay after OptimizeImpl mod end." << std::endl;
 
     // Get the updated function and new IRModule to build.
     // Instead of recreating the IRModule, we should look at the differences between this and the
@@ -431,15 +540,25 @@ class RelayBuildModule : public runtime::ModuleNode {
     // Generate code for the updated function.
     executor_codegen_ = MakeExecutorCodegen(executor_->name);
     executor_codegen_->Init(nullptr, config_->primitive_targets);
+    std::cout << "================= Codegen =================" << std::endl;
+    std::cout << "BuildRelay before Codegen mod:" << std::endl << PrettyPrint(func_module) << "BuildRelay before Codegen mod end." << std::endl;
     executor_codegen_->Codegen(func_module, func, mod_name);
+    std::cout << "=============== Codegen end ===============" << std::endl;
     executor_codegen_->UpdateOutput(&ret_);
     ret_.params = executor_codegen_->GetParams();
+    // std::cout << "[ZZDEBUG] BuildRelay ret_.params (only in graph executor):" << std::endl;
+    // int i = 0;
+    // for (const auto& kv : ret_.params) {
+    //   std::cout << "    ret_.params " << i << ": " << kv.first << std::endl;
+    //   i++;
+    // }
 
     auto lowered_funcs = executor_codegen_->GetIRModule();
 
     // No need to build for external functions.
     Target ext_dev("ext_dev");
     if (lowered_funcs.find(ext_dev) != lowered_funcs.end()) {
+      std::cout << "[ZZDEBUG] BuildRelay lowered_funcs.find(ext_dev) != lowered_funcs.end() (not in lenet)" << std::endl;
       lowered_funcs.Set(ext_dev, IRModule());
     }
 
@@ -458,15 +577,27 @@ class RelayBuildModule : public runtime::ModuleNode {
         ret_.mod = tvm::codegen::CSourceModuleCreate(";", "", Array<String>{});
       }
     } else {
+      std::cout << "================= TOLLVM =================" << std::endl;
+      std::cout << "[ZZDEBUG] lowered_funcs.size() = " << lowered_funcs.size() << std::endl;
+      for (const auto& kv : lowered_funcs) {
+        std::cout << "BuildRelay before TIRToRuntime lowered_funcs:" << std::endl << PrettyPrint(kv.second) << "BuildRelay before TIRToRuntime lowered_funcs end." << std::endl;
+      }
       ret_.mod = tvm::TIRToRuntime(lowered_funcs, host_target);
+      std::cout << "================= TOLLVM end =================" << std::endl;
     }
 
     auto ext_mods = executor_codegen_->GetExternalModules();
+    std::cout << "BuildRelay before CreateMetadataModule ret_.mod import size = " << ret_.mod->imports().size() << std::endl;
+    std::cout << "================= CreateMetadataModule =================" << std::endl;
     ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod, ext_mods, host_target,
                                                   runtime_, executor_,
                                                   executor_codegen_->GetExecutorCodegenMetadata());
+    std::cout << "================= CreateMetadataModule end =================" << std::endl;
+    // std::cout << "BuildRelay after CreateMetadataModule ret_.mod:" << std::endl << PrettyPrint(ret_.mod) << "BuildRelay after CreateMetadataModule ret_.mod end." << std::endl;
+    std::cout << "BuildRelay after CreateMetadataModule ret_.mod import size = " << ret_.mod->imports().size() << std::endl;
     // Remove external params which were stored in metadata module.
     for (tvm::runtime::Module mod : ext_mods) {
+      std::cout << "[ZZDEBUG] BuildRelay ext_mods not empty. (not in lenet)" << std::endl;
       auto pf_var = mod.GetFunction("get_const_vars");
       if (pf_var != nullptr) {
         Array<String> variables = pf_var();
@@ -479,6 +610,7 @@ class RelayBuildModule : public runtime::ModuleNode {
         }
       }
     }
+    std::cout << "================= BuildRelay end =================" << std::endl;
   }
 
  protected:
diff --git a/src/relay/backend/te_compiler.cc b/src/relay/backend/te_compiler.cc
index 816595474..a13c4fb70 100644
--- a/src/relay/backend/te_compiler.cc
+++ b/src/relay/backend/te_compiler.cc
@@ -147,8 +147,11 @@ class TECompilerImpl : public TECompilerNode {
   }
 
   // Lower the function.
-  CachedFunc Lower(const CCacheKey& key) {
-    return LowerInternal(key, global_var_supply_)->cached_func;
+  // CachedFunc Lower(const CCacheKey& key) {
+  //   return LowerInternal(key, global_var_supply_)->cached_func;
+  // }
+  CachedFunc Lower(const CCacheKey& key, bool use_default_sched = false) { // zz: add switch for default schedule
+    return LowerInternal(key, global_var_supply_, use_default_sched)->cached_func;
   }
 
   // TODO(gigiblender): Only to be called by the global TE compiler.
@@ -338,7 +341,8 @@ class TECompilerImpl : public TECompilerNode {
 
  private:
   // implement lowered func
-  CCacheValue LowerInternal(const CCacheKey& key, GlobalVarSupply global_var_supply) {
+  // CCacheValue LowerInternal(const CCacheKey& key, GlobalVarSupply global_var_supply) {
+  CCacheValue LowerInternal(const CCacheKey& key, GlobalVarSupply global_var_supply, bool use_default_sched = false) { // zz: add switch for default schedule
     VLOG(1) << "lowering:" << std::endl
             << PrettyPrint(key->source_func) << std::endl
             << "for target:" << std::endl
@@ -391,8 +395,17 @@ class TECompilerImpl : public TECompilerNode {
     With<Target> target_scope(key->target);
 
     ICHECK(!value->cached_func.defined());
-    value->cached_func =
-        PrimFuncFor(key->source_func, key->target, global_var_supply, constant_name_supply_);
+    // value->cached_func =
+    //     PrimFuncFor(key->source_func, key->target, global_var_supply, constant_name_supply_);
+    // zz: use_default_sched: switch between PrimFuncFor & PrimFuncForSimple
+    if (!use_default_sched) {
+      value->cached_func =
+          PrimFuncFor(key->source_func, key->target, global_var_supply, constant_name_supply_);
+    }
+    else {
+      value->cached_func =
+          PrimFuncForSimple(key->source_func, key->target, global_var_supply, constant_name_supply_);
+    }
 
     if (value->cached_func->prim_func.defined()) {
       VLOG(1) << "Lowering PrimFunc";
@@ -446,6 +459,23 @@ class TECompilerImpl : public TECompilerNode {
       VLOG(1) << "scheduling";
       IRModule scheduled_module = tvm::LowerSchedule(value->cached_func->schedule, all_args,
                                                      func_name, binds, global_var_supply);
+      // IRModule scheduled_module = tvm::LowerSchedule_ZZDEBUG(value->cached_func->schedule, all_args,
+      //                                                func_name, binds, global_var_supply);
+      // ---------------------------------------------------------------------------------
+      // zz: this cause bug
+      // ---------------------------------------------------------------------------------
+      // IRModule scheduled_module;
+      // if (!use_default_sched) {
+      //   IRModule scheduled_module = tvm::LowerSchedule(value->cached_func->schedule, all_args,
+      //                                                  func_name, binds, global_var_supply);
+      // }
+      // else {
+      //   IRModule scheduled_module = tvm::LowerSchedule(value->cached_func->schedule, all_args,
+      //                                                  func_name, binds, global_var_supply, true); // zz: simple_mode = true, to generate simple code
+      // }
+      // ---------------------------------------------------------------------------------
+      // zz: end - this cause bug
+      // ---------------------------------------------------------------------------------
       scheduled_module->Update(tir::transform::BindParams(all_consts)(scheduled_module));
       for (const auto& kv : scheduled_module->functions) {
         GlobalVar global_var = kv.first;
@@ -628,13 +658,19 @@ using AnalysisRemapping = std::unordered_map<Expr, Expr, ObjectHash, ObjectEqual
 class LowerTensorExprMutator : public DeviceAwareExprMutator {
  public:
   LowerTensorExprMutator(IRModule module, ProcessFn process_fn, CompilationConfig config,
-                         TECompiler compiler)
+                        //  TECompiler compiler)
+                        //  TECompiler compiler, bool use_default_sched = false) // zz: add switch for default schedule
+                         TECompiler compiler, bool use_default_sched = false, // zz: add switch for default schedule
+                         std::unordered_map<int, SubModItem> obfmod_functions_to_replace = {}) // zz: add obfuscated functions to replace
       : DeviceAwareExprMutator(module),
         module_(std::move(module)),
         process_fn_(std::move(process_fn)),
         config_(std::move(config)),
         compiler_(std::move(compiler)),
-        debug_op_(Op::Get("debug")) {}
+        // debug_op_(Op::Get("debug")) {}
+        debug_op_(Op::Get("debug")),
+        use_default_sched_(use_default_sched), // zz: add switch for default schedule
+        obfmod_functions_to_replace_(obfmod_functions_to_replace) {} // zz: add obfuscated functions to replace
 
   /*!
    *  \brief Returns the primitive function associated with \p expr, or nullptr if none.
@@ -677,6 +713,39 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
         // We have a regular call to a 'primitive' function (possibly with a 'Compiler' attribute).
         // We need to lower and rewrite the call.
         return GetRef<Function>(function_node);
+       /* ===========================================================================================================
+          * [zz] main logic: obfuscated functions is lowered already and stored in obfmod_functions_to_replace_
+          =========================================================================================================== */
+      } else if (function_node->attrs.GetAttr<Integer>(attr::kObf).defined()) { // zz: handle obfuscated functions
+        // std::cout << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive: expr.as FunctionNode kObf: " << std::endl << PrettyPrint(expr) << std::endl << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive: expr.as FunctionNode kObf end." << std::endl;
+        int obf_non_prim_func_id = function_node->attrs.GetAttr<Integer>(attr::kObf).value().IntValue();
+        if (obfmod_functions_to_replace_.find(obf_non_prim_func_id) != obfmod_functions_to_replace_.end()) {
+          // We have a regular call to an obfuscated function that needs to be replaced.
+          String obf_mod_name = obfmod_functions_to_replace_[obf_non_prim_func_id].first;
+          IRModule obf_mod = obfmod_functions_to_replace_[obf_non_prim_func_id].second;
+          for (auto omgv : obf_mod->GetGlobalVars()) {
+            if (omgv->name_hint.compare(obf_mod_name) == 0) {
+              BaseFunc obf_func = obf_mod->Lookup(omgv);
+              obf_prim_to_mod_[obf_func.get()] = obfmod_functions_to_replace_[obf_non_prim_func_id];
+              omgv->checked_type_ = obf_func->checked_type(); // zz: add checked_type_ from lowered Obf function
+              return obf_func;
+            }
+          }
+          ICHECK(false) << "Obfuscated lowered function not found (obf_mod obf_mod_name not found)";
+        }
+        else {
+          // std::cout << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive expr function_node ptr=" << std::hex << (uint64_t)(function_node) << std::endl;
+          // std::cout << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive expr:" << std::endl << PrettyPrint(expr) << std::endl << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive expr end." << std::endl;
+          // for (auto kv : obfmod_functions_to_replace_) {
+          //   std::cout << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive obfmod_functions_to_replace_ key=" << std::hex << (uint64_t)(kv.first) << std::endl;
+          //   std::cout << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive obfmod_functions_to_replace_ func:" << std::endl << PrettyPrint(GetRef<Function>(kv.first)) << std::endl << "[ZZDEBUG] LowerTensorExprMutator::ResolveToPrimitive obfmod_functions_to_replace_ func end." << std::endl;
+
+          // }
+          ICHECK(false) << "Obfuscated lowered function not found (obfmod_functions_to_replace_ key not found)";
+        }
+       /* ===========================================================================================================
+          * [zz] main logic end.
+          =========================================================================================================== */
       } else {
         // Not marked as primitive during partitioning or TVM fusion.
         return {};
@@ -819,6 +888,7 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 
   Expr DeviceAwareVisitExpr_(const FunctionNode* function_node) override {
     if (function_node->HasNonzeroAttr(attr::kPrimitive) ||
+        function_node->attrs.GetAttr<Integer>(attr::kObf).defined() || // zz: add Obf
         function_node->HasNonzeroAttr(attr::kExtern)) {
       // Nothing to lower inside primitive/external functions.
       return GetRef<Function>(function_node);
@@ -886,6 +956,50 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 
     // Case 4: If the function has already been lowered we just need to update the call.
     if (auto prim_func = primitive_func.as<tir::PrimFunc>()) {
+      /* ===========================================================================================================
+          * [zz] main logic: replace obfuscated functions with lowered return by ResolveToPrimitive
+         =========================================================================================================== */
+      // std::cout << "  (zz) LowerTensorExprMutator::DeviceAwareVisitExpr_ CallNode: case 4" << std::endl;
+      // std::cout << "  (zz) LowerTensorExprMutator::DeviceAwareVisitExpr_ CallNode: case 4 call_node->op: " << std::endl << PrettyPrint(call_node->op) << std::endl << "  (zz) LowerTensorExprMutator::DeviceAwareVisitExpr_ CallNode: case 4 call_node->op end." << std::endl;
+      if (obf_prim_to_mod_.find(primitive_func.get()) != obf_prim_to_mod_.end()) {
+        String obf_mod_name = obf_prim_to_mod_[primitive_func.get()].first;
+        IRModule obf_mod = obf_prim_to_mod_[primitive_func.get()].second;
+        for (auto omgv : obf_mod->GetGlobalVars()) {
+          if (omgv->name_hint.compare(obf_mod_name) == 0) {
+            Optional<Target> opt_target = primitive_func->GetAttr<Target>(tvm::attr::kTarget);
+            ICHECK(opt_target.defined());
+            // auto prim_fn_var = Downcast<GlobalVar>(call_node->op);
+            // Map<GlobalVar, BaseFunc> prim_fns = {{prim_fn_var, prim_func.value()}};
+            auto prim_fn_var = omgv;
+            // prim_fn_var->checked_type_ = primitive_func->checked_type();
+            // if (const auto* function_node = call_node->op.as<FunctionNode>()) {
+            //   Function func = GetRef<Function>(function_node);
+            //   prim_fn_var->checked_type_ = func->checked_type();
+            // }
+            // else {
+            //   ICHECK(false) << "Obfuscated lowered function not found (call_node->op not FunctionNode)";
+            // }
+            Map<GlobalVar, BaseFunc> prim_fns = {{prim_fn_var, prim_func.value()}};
+            // if (prim_fn_var->checked_type_.defined()) {
+            //   std::cout << "  (zz) LowerTensorExprMutator::DeviceAwareVisitExpr_ CallNode: prim_fn_var->checked_type_:" << prim_fn_var->checked_type_->GetTypeKey() << std::endl;
+            // }
+            // else {
+            //   std::cout << "  (zz) LowerTensorExprMutator::DeviceAwareVisitExpr_ CallNode: omgv->checked_type_ not defined" << std::endl;
+            // }
+            auto result = MakeLoweredCall(primitive_func, prim_fn_var, std::move(new_args), call_node->span,
+                                  opt_target.value(), prim_fns);
+            // std::cout << "  (zz) LowerTensorExprMutator::DeviceAwareVisitExpr_ CallNode: case 4 result:" << std::endl << PrettyPrint(result) << std::endl << "  (zz) LowerTensorExprMutator::DeviceAwareVisitExpr_ CallNode: case 4 result end." << std::endl;
+            // return MakeLoweredCall(primitive_func, omgv, std::move(new_args), call_node->span,
+            //                       opt_target.value(), prim_fns);
+            return result;
+          }
+        }
+        ICHECK(false) << "Obfuscated lowered function not found (obf_prim_to_mod_ key not found)";
+      }
+      /* ===========================================================================================================
+          * [zz] main logic end.
+         =========================================================================================================== */
+
       // Function should already be Target annotated by this point
       // but the TE Compiler metadata is still needed for the callback
       // TODO(Mousius) - Robustify this to not assume we're in the GlobalVar for Target Hooks
@@ -929,7 +1043,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
       // codegen.
       CCacheKey key(Downcast<Function>(primitive_func), target,
                     GetVirtualDevice(GetRef<Call>(call_node)));
-      CachedFunc cfunc = compiler_->Lower(key);
+      // CachedFunc cfunc = compiler_->Lower(key);
+      CachedFunc cfunc = compiler_->Lower(key, use_default_sched_); // zz: add switch for default schedule
       ICHECK(cfunc.defined());
       return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
                              call_node->span, target, cfunc->funcs->functions, cfunc->schedule);
@@ -949,20 +1064,28 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
   TECompiler compiler_;
   // Cache ops that need to be frequently used later to reduce lookup overhead.
   const Op& debug_op_;
+  bool use_default_sched_ = false; // zz: add switch for default schedule
+  std::unordered_map<int, SubModItem> obfmod_functions_to_replace_; // zz: add obfuscated functions to replace
+  std::unordered_map<const BaseFuncNode*, SubModItem> obf_prim_to_mod_; // zz: map lowered function back to its obf module
 };
 
-Pass LowerTensorExpr(TECompiler compiler, ProcessFn process_fn, CompilationConfig config) {
+// Pass LowerTensorExpr(TECompiler compiler, ProcessFn process_fn, CompilationConfig config) {
+Pass LowerTensorExpr(TECompiler compiler, ProcessFn process_fn, CompilationConfig config, bool use_default_sched, std::unordered_map<int, SubModItem> obfmod_functions_to_replace = {}) { // zz: add switch for default schedule & obfuscated functions to replace
   runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
       [=](Function func, IRModule module, PassContext ctx) {
-        LowerTensorExprMutator lower_te(module, process_fn, config, compiler);
+        // LowerTensorExprMutator lower_te(module, process_fn, config, compiler);
+        // std::cout << "[DEBUG] LowerTensorExpr: func passed to LowerTensorExprMutator: " << std::endl << func << "[DEBUG] LowerTensorExpr: func passed to LowerTensorExprMutator end." << std::endl;
+        LowerTensorExprMutator lower_te(module, process_fn, config, compiler, use_default_sched, obfmod_functions_to_replace); // zz: add switch for default schedule & obfuscated functions to replace
         return Downcast<Function>(lower_te.Mutate(func));
       };
   return CreateFunctionPass(pass_func, 0, "LowerTensorExpr", {});
 }
 
 backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod, const CompilationConfig& config,
-                                              Map<Expr, backend::StorageInfo> storage_info_map) {
-  Function func = Downcast<Function>(mod->Lookup("main"));
+                                              // Map<Expr, backend::StorageInfo> storage_info_map) {
+                                              Map<Expr, backend::StorageInfo> storage_info_map, String main_func_name) { // zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+  // Function func = Downcast<Function>(mod->Lookup("main"));
+  Function func = Downcast<Function>(mod->Lookup(main_func_name)); // zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
 
   VLOG_CONTEXT << "UpdateMainWorkspaceSize";
   VLOG(1) << "calculating FunctionInfo for main:" << std::endl << PrettyPrint(func);
@@ -1185,7 +1308,11 @@ void UpdateFunctionMetadata(BaseFunc func,
 
 /*! \brief Main lowering driving. */
 IRModule LowerTE(const IRModule& module, const String& module_name, ProcessFn process_fn,
-                 CompilationConfig config) {
+                //  CompilationConfig config) {
+                //  CompilationConfig config, bool use_default_sched = false) { // zz: add switch to use default_sched
+                 CompilationConfig config, bool use_default_sched = false, // zz: add switch to use default_sched
+                 std::unordered_map<int, SubModItem> obfmod_functions_to_replace = {}) { // zz: add obfmod_functions_to_replace
+  // std::cout << "[ZZDEBUG] LowerTE() module: " << std::endl << PrettyPrint(module) << std::endl << "[ZZDEBUG] LowerTE() module end." << std::endl;
   TECompiler compiler(module, module_name);
 
   // TODO(mbs): This is all unnecessarily convoluted. Better would be to accumulate the rewritten
@@ -1201,7 +1328,8 @@ IRModule LowerTE(const IRModule& module, const String& module_name, ProcessFn pr
   //  - Calls to functions tagged with "Primitive" are compiled to PrimFuncs, and calls updated
   //    (using call_lowered convention).
   IRModule updated_module =
-      LowerTensorExpr(compiler, std::move(process_fn), std::move(config))(module);
+      // LowerTensorExpr(compiler, std::move(process_fn), std::move(config))(module);
+      LowerTensorExpr(compiler, std::move(process_fn), std::move(config), use_default_sched, obfmod_functions_to_replace)(module); // zz: add default_sched switch & obfmod_functions_to_replace
 
   // The Functions tagged with "Compiler" are now residing in the cache ready to be
   // compiled by LowerExternalFunctions. However we still need a record of them in the
@@ -1223,6 +1351,14 @@ IRModule LowerTE(const IRModule& module, const String& module_name, ProcessFn pr
     updated_module->Add(kv.first, kv.second);
   }
 
+  // zz: add obfuscated functions module
+  for (const auto& fun_mod : obfmod_functions_to_replace) {
+    auto obf_mod = fun_mod.second.second;
+    for (const auto& kv : obf_mod->functions) {
+      updated_module->Add(kv.first, kv.second);
+    }
+  }
+
   // Invoke external codegen for all Functions in the cache tagged with "Compiler", and
   // annotate the module with the resulting runtime modules.
   // TODO(mbs): runtime modules should be first class rather than attributes.
@@ -1313,6 +1449,19 @@ Pass LowerTE(String module_name, CompilationConfig complilation_config, ProcessF
        tvm::tir::transform::ExtractPrimFuncConstants()});
 }
 
+// zz: add switch for default schedule
+Pass LowerTE_Simple(String module_name, CompilationConfig complilation_config, ProcessFn process_fn, std::unordered_map<int, SubModItem> obfmod_functions_to_replace) {
+  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func = [=](IRModule module,
+                                                                            PassContext ctx) {
+    return LowerTE(module, module_name, process_fn, complilation_config, true, obfmod_functions_to_replace); // zz: use default_sched
+  };
+
+  return tvm::transform::Sequential(
+      {tvm::relay::transform::RelayToTIRTargetHook(complilation_config),
+       tvm::transform::CreateModulePass(pass_func, 0, "LowerTE", {"InferType"}), InferType(),
+       tvm::tir::transform::ExtractPrimFuncConstants()}, "LowerTE_ZZDEBUG"); // zz: use LowerTE_ZZDEBUG as Sequential name for debug purpose (check src/ir/transform.cc)
+}
+
 TVM_REGISTER_GLOBAL("relay.tec.LowerTE")
     .set_body_typed([](String module_name, CompilationConfig compilation_config) {
       return LowerTE(std::move(module_name), std::move(compilation_config));
diff --git a/src/relay/backend/te_compiler.h b/src/relay/backend/te_compiler.h
index f2ba84014..3213f964d 100644
--- a/src/relay/backend/te_compiler.h
+++ b/src/relay/backend/te_compiler.h
@@ -54,6 +54,8 @@
 #include "./te_compiler_cache.h"
 #include "./utils.h"
 
+using SubModItem = std::pair<tvm::runtime::String, tvm::IRModule>;
+
 namespace tvm {
 namespace relay {
 namespace tec {
@@ -73,7 +75,8 @@ class TECompilerNode : public Object {
    * \param key The key to the cached function.
    * \return The result.
    */
-  virtual CachedFunc Lower(const CCacheKey& key) = 0;
+  // virtual CachedFunc Lower(const CCacheKey& key) = 0;
+  virtual CachedFunc Lower(const CCacheKey& key, bool use_default_sched = false) = 0;
 
   /*!
    * \brief Get lowered result.
@@ -163,7 +166,8 @@ void UpdateFunctionMetadata(BaseFunc relay_func,
  * \return function_infos Function info for each function in the module
  */
 backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod, const CompilationConfig& config,
-                                              Map<Expr, backend::StorageInfo> storage_info_map);
+                                              // Map<Expr, backend::StorageInfo> storage_info_map);
+                                              Map<Expr, backend::StorageInfo> storage_info_map, String main_func_name="main"); // zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
 
 /*! \brief Returns all the global \p PrimFunc functions in \p mod, but separated into an \p IRModule
  * per \p Target.
@@ -189,6 +193,9 @@ inline void DefaultProcessFn(BaseFunc) {}
  */
 transform::Pass LowerTE(String module_name, CompilationConfig config,
                         ProcessFn process_fn = DefaultProcessFn);
+// zz: add switch for default schedule
+transform::Pass LowerTE_Simple(String module_name, CompilationConfig config,
+                        ProcessFn process_fn = DefaultProcessFn, std::unordered_map<int, SubModItem> obfmod_functions_to_replace = {});
 
 }  // namespace tec
 }  // namespace relay
diff --git a/src/relay/backend/te_compiler_cache.cc b/src/relay/backend/te_compiler_cache.cc
index b747855bf..d4b17e57a 100644
--- a/src/relay/backend/te_compiler_cache.cc
+++ b/src/relay/backend/te_compiler_cache.cc
@@ -711,6 +711,315 @@ class ScheduleBuilder : public ExprVisitor {
                       lower_te_compute.constant_tensors_);
   }
 
+  // zz: use default default_sched = GenericFunc::Get("schedule_injective") for simplicity, ref Create()
+  CachedFunc CreateSimple(const Function& relay_func, GlobalVarSupply global_var_supply,
+                    NameSupply constant_name_supply) {
+    LowerToTECompute lower_te_compute(target_, constant_name_supply);
+    // std::cout << "[ZZDEBUG] LowerToTECompute for relay_func (fused op functions after FuseOps):" << std::endl << PrettyPrint(relay_func) << std::endl;
+    Array<te::Tensor> tensor_outs = lower_te_compute.Lower(relay_func);
+    // std::cout << "[ZZDEBUG] lower_te_compute.Lower results:" << std::endl;
+    // std::cout << "  (zz) CreateSimple tensor_outs.size() = " << tensor_outs.size() << std::endl;
+    // for (auto tensor : tensor_outs) {
+    //   std::cout << "    (tensor_outs ele) tensor = " << PrettyPrint(tensor) << std::endl;
+    //   std::cout << "    (tensor_outs ele) tensor->op (carries the computation from input to output) = " << PrettyPrint(tensor->op) << std::endl;
+    //   if (tensor->op.as<te::ComputeOpNode>()) {
+    //     // std::cout << "    (tensor_outs ele) tensor->op->body = " << std::endl << PrettyPrint(tensor->op.as<te::ComputeOpNode>()->body) << std::endl;
+    //     for (auto body_ele : tensor->op.as<te::ComputeOpNode>()->body) {
+    //       std::cout << "      (tensor ele) tensor->op->body ele type = " << body_ele->GetTypeKey() << std::endl;
+    //       std::cout << "      (tensor ele) tensor->op->body ele = " << PrettyPrint(body_ele) << std::endl;
+    //     }
+    //   }
+    // }
+    Array<te::Tensor> fn_inputs = lower_te_compute.fn_inputs_;
+    VisitExpr(relay_func->body);
+
+    // TODO(mbs): This should be the definitive global by which the PrimFunc is known and
+    // no other GlobalVar ctors should appear inside the lowering machinery.
+    auto prim_fn_var = global_var_supply->FreshGlobal(lower_te_compute.candidate_name_);
+    prim_fn_var->checked_type_ = relay_func->checked_type();
+
+    // Fusion over tupled results may leave identity relationships
+    // between inputs and outputs, copy identity output tensors,
+    // since tir lowering do not support aliasing output to input buffer.
+    for (size_t i = 0; i < tensor_outs.size(); ++i) {
+      if (tensor_outs[i]->op.as<te::PlaceholderOpNode>()) {
+        tensor_outs.Set(i, topi::identity(tensor_outs[i]));
+      }
+    }
+
+    te::Schedule schedule{nullptr};
+    tir::PrimFunc prim_func{nullptr};
+    // No need to register schedule for device copy op.
+    if (anchor_attrs_.as<DeviceCopyAttrs>() == nullptr) {
+      if (use_auto_scheduler_) {
+        const auto* fauto_schedule =
+            runtime::Registry::Get("auto_scheduler.relay_integration.auto_schedule_topi_compute");
+        ICHECK(fauto_schedule != nullptr)
+            << "auto_scheduler.relay_integration.auto_schedule_topi_compute is not registered";
+        ObjectRef obj = (*fauto_schedule)(prim_fn_var->name_hint, tensor_outs);
+        if (obj.defined()) {
+          schedule = Downcast<te::Schedule>(obj);
+        }
+      }
+      if (database_) {
+        using tvm::meta_schedule::TuningRecord;
+        using tvm::tir::IndexMap;
+        using tvm::tir::Instruction;
+        using tvm::tir::InstructionKind;
+        using tvm::tir::PrimFunc;
+        using tvm::tir::Schedule;
+        backend::FTECompilerTIRConverter tir_converter = backend::GetTIRConverter();
+        Array<te::Tensor> te_args = Concat(fn_inputs, tensor_outs);
+        Array<runtime::NDArray> constants;
+        for (auto [const_node, te_tensor] : lower_te_compute.constant_tensors_) {
+          te_args.push_back(te_tensor);
+          constants.push_back(const_node->data);
+        }
+        if (Optional<PrimFunc> f = tir_converter(te_args, constants)) {
+          IRModule query_mod = backend::PrimFuncToIRModule(f.value());
+          if (Optional<TuningRecord> opt_record = database_.value()->QueryTuningRecord(
+                  /*mod=*/query_mod,
+                  /*target=*/target_,
+                  /*workload_name=*/prim_fn_var->name_hint)) {
+            LayoutFreeConstantCollector const_collector;
+            const_collector(f.value()->body);
+
+            static InstructionKind kind_transform_layout = InstructionKind::Get("TransformLayout");
+            TuningRecord record = opt_record.value();
+            for (const Instruction& inst : record->trace->insts) {
+              if (inst->kind.same_as(kind_transform_layout)) {
+                ICHECK_EQ(inst->inputs.size(), 2);
+                auto index_map = Downcast<IndexMap>(inst->inputs[1]);
+
+                if (!const_collector.constants.empty()) {
+                  // In this case, RewriteLayout is acting on an AllocateConst node.
+                  // After tuning, we reach this code path twice: First by
+                  // the Relay MetaScheduleLayoutRewrite pass, and next by the final
+                  // compilation (Relay to TE schedule lowering).
+                  //
+                  // Due to Relay MetaScheduleLayoutRewrite and FoldConstant passes,
+                  // the Relay subgraph for which we query the database during the
+                  // final compilation has its weight tensor transformed according to
+                  // the index map, determined during tuning. For example,
+                  //
+                  // fn (%p0: Tensor[(1, 56, 56, 64), float32]) {
+                  //   %0 = nn.conv2d(%p0, meta[relay.Constant][0],
+                  //                       /*ty=Tensor[(4, 2, 2, 3, 3, 32, 8), float32]*/, ...);
+                  //   add(%0, meta[relay.Constant][1])
+                  // }
+                  //
+                  // Note that the database does not have an entry corresponding to such subgraphs,
+                  // since an input subgraph to the tuning system always has its weight tensor in
+                  // the original layout, e.g.
+                  //
+                  // fn (%p0: Tensor[(1, 56, 56, 64), float32]) {
+                  //   %0 = nn.conv2d(%p0, meta[relay.Constant][0],
+                  //                       /*ty=Tensor[(3, 3, 64, 64), float32]*/, ...);
+                  //   add(%0, meta[relay.Constant][1])
+                  // }
+                  //
+                  // Thus, in both of the two cases where we reach this code path, we need careful
+                  // logic to make sure that (1) the database lookup during the final compilation
+                  // succeeds and (2) the application of a schedule trace is well defined.
+
+                  ICHECK(const_collector.constants.size() == 1)
+                      << "Only one layout-free constant is supported by RewriteLayout for now";
+                  auto constant = const_collector.constants[0];
+
+                  auto is_constant_transformed = [index_map](runtime::NDArray c) {
+                    if (c.Shape().size() != index_map->initial_indices.size()) {
+                      return true;
+                    }
+                    size_t src_size_1d = 1;
+                    Array<PrimExpr> orig_shape;
+                    for (size_t i = 0; i < c.Shape().size(); ++i) {
+                      src_size_1d *= c->shape[i];
+                      orig_shape.push_back(PrimExpr(static_cast<int>((c->shape[i]))));
+                    }
+                    arith::Analyzer analyzer;
+                    auto dst_shape = index_map->MapShape(orig_shape, &analyzer);
+                    std::vector<int64_t> dst_shape_int;
+                    size_t dst_size_1d = 1;
+                    for (size_t i = 0; i < dst_shape.size(); ++i) {
+                      dst_size_1d *= dst_shape[i].as<IntImmNode>()->value;
+                    }
+                    return src_size_1d != dst_size_1d;
+                  };
+
+                  if (!is_constant_transformed(constant)) {
+                    // This is the first case, reached during the MetaScheduleLayoutRewrite pass.
+                    //
+                    // A layout-free constant having the same rank as an input to the index map
+                    // is assumed to be transformed by this index map.
+                    // TODO(masahi): If there are multiple layout-free constants in one
+                    // TIR mod (e.g. conv2d -> conv2d fusion), this assumption does not hold.
+                    // We need to determine which constant the given index map acts on.
+                    //
+                    // We know that, during the final compilation, we will query the database
+                    // for a subgraph that the tuner has never seen. We workaround this problem
+                    // by adding a dummy entry to the database. The dummy entry is carefully
+                    // constructed so that the lookup during the final compilation would succeed.
+                    runtime::NDArray rewritten_constant = index_map->MapNDArray(constant);
+                    auto f_dummy = AllocateConstReplaceConstant::Rewrite(
+                        f.value(), {{constant, rewritten_constant}});
+                    auto workload_dummy =
+                        database_.value()->CommitWorkload(backend::PrimFuncToIRModule(f_dummy));
+                    TuningRecord rec_dummy(record->trace, workload_dummy, record->run_secs,
+                                           record->target, record->args_info);
+                    database_.value()->CommitTuningRecord(rec_dummy);
+                  } else {
+                    // The constant is already transformed, so this is the second case, reached
+                    // during the final compilation.
+                    //
+                    // The schedule trace is supposed to be applied to the weight in its original
+                    // layout. But as explained above, the Relay subgraph we get in this case
+                    // has its weight tensor transformed according to the corresponding index map.
+                    // So effectively, we undo the layout transformation on the weight to restore
+                    // the original PrimFunc that the schedule trace is supposed to act on.
+                    ICHECK(index_map->inverse_index_map);
+                    auto inverse_map = Downcast<IndexMap>(index_map->inverse_index_map.value());
+                    ICHECK(constant.Shape().size() == inverse_map->initial_indices.size());
+                    runtime::NDArray orig_constant = inverse_map->MapNDArray(constant);
+                    auto f_ = AllocateConstReplaceConstant::Rewrite(f.value(),
+                                                                    {{constant, orig_constant}});
+                    query_mod = backend::PrimFuncToIRModule(f_);
+                  }
+                }
+                MetaScheduleLayoutRewriter::LayoutQueuePush(index_map);
+              }
+            }
+
+            Schedule sch = Schedule::Traced(query_mod, /*seed=*/-1, /*debug_mask=*/0,
+                                            tir::ScheduleErrorRenderLevel::kDetail);
+
+            if (!mod_eq_structural_->Equal(query_mod, opt_record.value()->workload->mod)) {
+              // When the database lookup succeeds while structural equality check fails,
+              // it implies that the anchor block based equality has been used during tuning.
+              // The trace in the record cannot directly be applied to this query module.
+              meta_schedule::ScheduleUsingAnchorTrace(sch, record->trace, target_);
+            } else {
+              record->trace->ApplyToSchedule(sch, /*remove_postproc=*/false);
+            }
+
+            IRModule mod = sch->mod();
+            ICHECK_EQ(mod->functions.size(), 1);
+            mod = tir::transform::RemoveWeightLayoutRewriteBlock(/*skip_ndarray_rewrite*/ false)(
+                std::move(mod));
+            prim_func = Downcast<PrimFunc>(mod->Lookup("main"));
+            // Need to copy attrs from relay function over to prim func. Most notably the structural
+            // hash.
+            prim_func = WithAttrs(prim_func, relay_func->attrs->dict);
+          } else {
+            int dispatch = backend::UseMetaScheduleDispatch();
+            // (dispatch & 2): controls whether to print TVMScript for missing TIR
+            // (dispatch & 4): controls whether to raise fatal errors for missing TIR
+            if (dispatch & 2) {
+              LOG(WARNING) << "Cannot find workload: " << prim_fn_var->name_hint << "\n"
+                           << f.value();
+            } else {
+              LOG(WARNING) << "Cannot find workload: " << prim_fn_var->name_hint;
+            }
+            if (dispatch & 4) {
+              LOG(FATAL);
+            }
+          }
+        }
+      }
+      // Use TOPI schedule if user specified, or the function has no auto_scheduler schedule.
+      if (!schedule.defined() && !prim_func.defined()) {
+        // zz: just use default_sched for simplicity
+        // if (anchor_op_.defined()) {
+        //   std::cout << "[ZZDEBUG] ScheduleBuilder::Create: anchor_impl\n";
+        //   auto anchor_impl = lower_te_compute.op_implementations_.find(anchor_op_.operator->());
+        //   ICHECK(anchor_impl != lower_te_compute.op_implementations_.end());
+        //   schedule = anchor_impl->second.Schedule(anchor_attrs_, tensor_outs, target_);
+        // } else {
+        // std::cout << "[ZZDEBUG] ScheduleBuilder::Create: default_sched\n";
+        auto default_sched = GenericFunc::Get("schedule_injective");
+        ICHECK(default_sched.defined()) << "schedule_injective not registered for " << target_;
+        With<Target> tctx(target_);
+        schedule = default_sched(tensor_outs);
+        // }
+
+        // zz: do reordering for conv2d_NCHWc
+        transform::PassContext pass_ctx = transform::PassContext::Current();
+        bool enable_compute_reorder = pass_ctx->GetConfig<Bool>("dnnobf.compute_reorder", Bool(false)).value();
+        if (enable_compute_reorder) {
+          // std::cout << "[ZZDEBUG] default_sched schedule stage_map.size=" << schedule->stage_map.size() << std::endl;
+          for (auto kv : schedule->stage_map) {
+            // std::cout << "  stage_map[" << kv.first << "] = " << kv.second << std::endl;
+            // std::cout << "    (zz) key: " << kv.first << std::endl;
+            // std::cout << "    (zz) value: " << kv.second << std::endl;
+            // // print all_iter_vars
+            // std::cout << "    (zz) value name: " << kv.second->op->name << std::endl;
+            // std::cout << "      (zz) all_iter_vars.size() = " << kv.second->all_iter_vars.size() << std::endl;
+            // for (auto iter_var : kv.second->all_iter_vars) {
+            //   std::cout << "        (zz) iter_var = " << PrettyPrint(iter_var) << std::endl;
+            // }
+
+            if (kv.second->op->name == "conv2d_NCHWc") {
+              Array<String> candidate_conv2d_NCHWc_iter_vars = {"n", "oc_chunk", "oh", "ow", "oc_block", "ic", "kh", "kw"};
+              // reorder {"n", "oc_chunk", "oh", "ow", "oc_block", "ic", "kh", "kw"} to {"n", "oc_chunk", "ic", "kh", "kw", "oh", "ow", "oc_block"}
+              ICHECK(kv.second->all_iter_vars.size() == 8);
+              std::vector<IterVar> new_iter_vars;
+              for (auto iter_var : kv.second->all_iter_vars) {
+                if (std::find(candidate_conv2d_NCHWc_iter_vars.begin(), candidate_conv2d_NCHWc_iter_vars.end(), iter_var->var->name_hint) != candidate_conv2d_NCHWc_iter_vars.end()) {
+                  new_iter_vars.push_back(iter_var);
+                }
+              }
+              ICHECK(new_iter_vars.size() == 8);
+              schedule[kv.first].reorder({new_iter_vars[0], new_iter_vars[5], new_iter_vars[6], new_iter_vars[7], new_iter_vars[1], new_iter_vars[2], new_iter_vars[3], new_iter_vars[4]});
+              std::cout << "[DNNOBF] reorder conv2d_NCHWc iter_vars from {n, oc_chunk, oh, ow, oc_block, ic, kh, kw} to {n, oc_chunk, ic, kh, kw, oh, ow, oc_block}" << std::endl;
+            }
+
+            if (kv.second->op->name == "T_batch_matmul_NT") {
+              Array<String> candidate_T_batch_matmul_NT_iter_vars = {"b", "i", "j", "k"};
+              // reorder {"b", "i", "j", "k"} to {"k", "b", "i", "j"}
+              ICHECK(kv.second->all_iter_vars.size() == 4);
+              std::vector<IterVar> new_iter_vars;
+              for (auto iter_var : kv.second->all_iter_vars) {
+                if (std::find(candidate_T_batch_matmul_NT_iter_vars.begin(), candidate_T_batch_matmul_NT_iter_vars.end(), iter_var->var->name_hint) != candidate_T_batch_matmul_NT_iter_vars.end()) {
+                  new_iter_vars.push_back(iter_var);
+                }
+              }
+              ICHECK(new_iter_vars.size() == 4);
+              schedule[kv.first].reorder({new_iter_vars[3], new_iter_vars[0], new_iter_vars[1], new_iter_vars[2]});
+              std::cout << "[DNNOBF] reorder T_batch_matmul_NT iter_vars from {b, i, j, k} to {k, b, i, j}" << std::endl;
+            }
+
+            // for fused_nn_dense_constant, see fastext for example
+            if (kv.second->op->name == "compute") {
+              Array<String> candidate_fused_nn_dense_constant_vars = {"y", "x", "k"};
+              // reorder {"y", "x", "k"} to {"y", "k", "x"}
+              ICHECK(kv.second->all_iter_vars.size() == 3);
+              std::vector<IterVar> new_iter_vars;
+              for (auto iter_var : kv.second->all_iter_vars) {
+                if (std::find(candidate_fused_nn_dense_constant_vars.begin(), candidate_fused_nn_dense_constant_vars.end(), iter_var->var->name_hint) != candidate_fused_nn_dense_constant_vars.end()) {
+                  new_iter_vars.push_back(iter_var);
+                }
+              }
+              ICHECK(new_iter_vars.size() == 3);
+              schedule[kv.first].reorder({new_iter_vars[0], new_iter_vars[2], new_iter_vars[1]});
+              std::cout << "[DNNOBF] reorder compute iter_vars from {y, x, k} to {y, k, x}" << std::endl;
+            }
+          }
+        }
+
+      }
+      if (schedule.defined()) {
+        for (const auto& scalar : lower_te_compute.scalars_) {
+          if (schedule->Contain(scalar)) {
+            schedule[scalar].compute_inline();
+          }
+        }
+      }
+    }
+
+    IRModule funcs = IRModule(Map<GlobalVar, BaseFunc>({}));
+    return CachedFunc(target_, prim_fn_var, fn_inputs, tensor_outs, schedule, prim_func, {}, funcs,
+                      lower_te_compute.constant_tensors_);
+  }
+
   void VisitExpr_(const CallNode* call_node) final {
     static auto fpattern = Op::GetAttrMap<TOpPattern>("TOpPattern");
 
@@ -722,11 +1031,12 @@ class ScheduleBuilder : public ExprVisitor {
     }
 
     int op_pattern = fpattern[op];
-    if (!use_auto_scheduler_ && !database_.defined() && op_pattern >= kCommReduce) {
-      ICHECK(!anchor_op_.defined() || anchor_op_pattern_ < kCommReduce)
-          << "Cannot apply TOPI schedule to a primitive function with two complicated ops"
-          << " anchor=" << anchor_op_ << " current=" << op;
-    }
+    // zz: remove the check for complicated ops fusion
+    // if (!use_auto_scheduler_ && !database_.defined() && op_pattern >= kCommReduce) {
+    //   ICHECK(!anchor_op_.defined() || anchor_op_pattern_ < kCommReduce)
+    //       << "Cannot apply TOPI schedule to a primitive function with two complicated ops"
+    //       << " anchor=" << anchor_op_ << " current=" << op;
+    // }
     if (op_pattern >= anchor_op_pattern_) {
       anchor_op_ = op;
       anchor_attrs_ = call_node->attrs;
@@ -756,6 +1066,13 @@ CachedFunc PrimFuncFor(const Function& source_func, const Target& target,
   return ScheduleBuilder(target).Create(source_func, global_var_supply, constant_name_supply);
 }
 
+// zz: use CreateSimple for simplicity
+CachedFunc PrimFuncForSimple(const Function& source_func, const Target& target,
+                       GlobalVarSupply global_var_supply, NameSupply constant_name_supply) {
+  return ScheduleBuilder(target).CreateSimple(source_func, global_var_supply, constant_name_supply);
+}
+
+
 // Creates shape function from functor.
 class MakeShapeFunc : public backend::MemoizedExprTranslator<Array<te::Tensor>> {
  public:
diff --git a/src/relay/backend/te_compiler_cache.h b/src/relay/backend/te_compiler_cache.h
index 76939a923..ed6fd3076 100644
--- a/src/relay/backend/te_compiler_cache.h
+++ b/src/relay/backend/te_compiler_cache.h
@@ -248,6 +248,10 @@ std::pair<Optional<tir::PrimFunc>, std::string> LowerToPrimFunc(const Function&
 CachedFunc PrimFuncFor(const Function& source_func, const Target& target,
                        GlobalVarSupply global_var_supply, NameSupply constant_name_supply);
 
+// zz: use default default_sched for simplicity
+CachedFunc PrimFuncForSimple(const Function& source_func, const Target& target,
+                       GlobalVarSupply global_var_supply, NameSupply constant_name_supply);
+
 /*! \brief A specialization of PrimFuncFor, meant to be used when the names of constants do not
  * matter. */
 inline CachedFunc PrimFuncFor(const Function& source_func, const Target& target) {
@@ -268,10 +272,11 @@ inline size_t CCacheKeyNode::Hash() const {
 }
 
 inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
-  if (Hash() != other->Hash()) return false;
-  return this->target->str() == other->target->str() &&
-         this->virtual_device == other->virtual_device &&
-         tvm::StructuralEqual()(this->source_func, other->source_func);
+  // if (Hash() != other->Hash()) return false;
+  // return this->target->str() == other->target->str() &&
+  //        this->virtual_device == other->virtual_device &&
+  //        tvm::StructuralEqual()(this->source_func, other->source_func);
+  return false;  // zz: disable cache for transformer models to avoid bug
 }
 
 }  // namespace tec
diff --git a/src/relay/backend/utils.cc b/src/relay/backend/utils.cc
index f009bda9c..65f1c2990 100644
--- a/src/relay/backend/utils.cc
+++ b/src/relay/backend/utils.cc
@@ -268,14 +268,20 @@ Array<Pass> GetPassPrefix(bool is_homogeneous, bool is_vm) {
   pass_seqs.push_back(transform::CanonicalizeOps());
   pass_seqs.push_back(transform::FlattenAtrousConv());
 
-  // Alter layout transformation is currently only applied to homogeneous execution.
-  if (is_homogeneous) {
-    if (!is_vm) {
-      pass_seqs.push_back(transform::InferType());
-    }
-    pass_seqs.push_back(transform::AlterOpLayout());
-    pass_seqs.push_back(transform::SimplifyExprPostAlterOp());
-  }
+  // ====================================================
+  // zz: disable this for simplicity
+  // ====================================================
+  // // Alter layout transformation is currently only applied to homogeneous execution.
+  // if (is_homogeneous) {
+  //   if (!is_vm) {
+  //     pass_seqs.push_back(transform::InferType());
+  //   }
+  //   pass_seqs.push_back(transform::AlterOpLayout());
+  //   pass_seqs.push_back(transform::SimplifyExprPostAlterOp());
+  // }
+  // ====================================================
+  // zz: disable this for simplicity end
+  // ====================================================
 
   // Fast math optimizations.
   pass_seqs.push_back(transform::FastMath());
diff --git a/src/relay/ir/expr.cc b/src/relay/ir/expr.cc
index 062d9206c..6f660e849 100644
--- a/src/relay/ir/expr.cc
+++ b/src/relay/ir/expr.cc
@@ -230,7 +230,8 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
       p->stream << ")";
     });
 
-Call::Call(Expr op, Array<Expr> args, Attrs attrs, Array<Type> type_args, Span span) {
+// Call::Call(Expr op, Array<Expr> args, Attrs attrs, Array<Type> type_args, Span span) {
+Call::Call(Expr op, Array<Expr> args, Attrs attrs, Array<Type> type_args, Span span, DictAttrs obf_attrs) { // zz: add obf_attrs
   ObjectPtr<CallNode> n = make_object<CallNode>();
   n->op = std::move(op);
   n->args = std::move(args);
@@ -238,22 +239,26 @@ Call::Call(Expr op, Array<Expr> args, Attrs attrs, Array<Type> type_args, Span s
   n->type_args = std::move(type_args);
   n->virtual_device_ = VirtualDevice::FullyUnconstrained();
   n->span = std::move(span);
+  n->obf_attrs = std::move(obf_attrs); // zz: add obf_attrs
   data_ = std::move(n);
 }
 
 Call WithFields(Call call, Optional<Expr> opt_op, Optional<Array<Expr>> opt_args,
                 Optional<Attrs> opt_attrs, Optional<Array<Type>> opt_type_args,
-                Optional<VirtualDevice> opt_virtual_device, Optional<Span> opt_span) {
+                // Optional<VirtualDevice> opt_virtual_device, Optional<Span> opt_span) {
+                Optional<VirtualDevice> opt_virtual_device, Optional<Span> opt_span, Optional<DictAttrs> opt_obf_attrs) { // zz: add obf_attrs
   // Collect new values for fields.
   Expr op = opt_op.value_or(call->op);
   Array<Expr> args = opt_args.value_or(call->args);
   Attrs attrs = opt_attrs.value_or(call->attrs);
+  DictAttrs obf_attrs = opt_obf_attrs.value_or(call->obf_attrs); // zz: add obf_attrs
   Array<Type> type_args = opt_type_args.value_or(call->type_args);
   VirtualDevice virtual_device = opt_virtual_device.value_or(call->virtual_device());
   Span span = opt_span.value_or(call->span);
 
   // Check if anything changed.
   bool unchanged = op.same_as(call->op) && attrs.same_as(call->attrs) &&
+                   obf_attrs.same_as(call->obf_attrs) && // zz: add obf_attrs
                    virtual_device.same_as(call->virtual_device()) && span.same_as(call->span);
   if (unchanged) {
     if (args.size() == call->args.size()) {
@@ -283,10 +288,24 @@ Call WithFields(Call call, Optional<Expr> opt_op, Optional<Array<Expr>> opt_args
     cow_call_node->type_args = type_args;
     cow_call_node->virtual_device_ = virtual_device;
     cow_call_node->span = span;
+    cow_call_node->obf_attrs = obf_attrs;
   }
   return call;
 }
 
+// zz: for obfuscation, ref: include/tvm/ir/attrs.h WithAttr
+Call WithObfAttr(Call input, const std::string& attr_key, ObjectRef attr_value) {
+  CallNode* node = input.CopyOnWrite();
+  if (node->obf_attrs.defined()) {
+    node->obf_attrs.CopyOnWrite()->dict.Set(attr_key, attr_value);
+  } else {
+    Map<String, ObjectRef> dict = {{attr_key, attr_value}};
+    node->obf_attrs = DictAttrs(dict);
+  }
+  return input;
+}
+
+
 TVM_REGISTER_NODE_TYPE(CallNode);
 
 TVM_REGISTER_GLOBAL("relay.ir.Call")
@@ -308,6 +327,16 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
                 << node->type_args << ")";
     });
 
+// zz: for obfuscation
+TVM_REGISTER_GLOBAL("relay.ir.CallWithObfAttr")
+    .set_body_typed([](Call call, const std::string& attr_key, ObjectRef attr_value) {
+      return WithObfAttr(call, attr_key, attr_value);
+    });
+// ref: src/ir/module.cc GetAttr ir.Module_GetAttr
+TVM_REGISTER_GLOBAL("relay.ir.CallGetObfAttr").set_body_typed([](Call call, String key) -> ObjectRef {
+  return GetObfAttr<ObjectRef>(call, key);
+});
+
 Let::Let(Var var, Expr value, Expr body, Span span) {
   ObjectPtr<LetNode> n = make_object<LetNode>();
   n->var = std::move(var);
diff --git a/src/relay/printer/relay_text_printer.cc b/src/relay/printer/relay_text_printer.cc
index 618e8fe13..f7c40d898 100644
--- a/src/relay/printer/relay_text_printer.cc
+++ b/src/relay/printer/relay_text_printer.cc
@@ -510,6 +510,14 @@ Doc RelayTextPrinter::VisitExpr_(const CallNode* op) {
   for (const Doc& d : PrintCallAttrs(op->attrs, op->op)) {
     args.push_back(d);
   }
+  // zz: print obf_attrs
+  if (op->obf_attrs.defined()) {
+    // args.emplace_back(PrintDictAttrs(op->obf_attrs));
+    for (const Doc& d : PrintDictAttrs(op->obf_attrs)) {
+      args.push_back(d);
+    }
+  }
+
   const auto* cons_node = op->op.as<ConstructorNode>();
   if (cons_node) {
     doc << cons_node->name_hint;
diff --git a/src/relay/transforms/device_aware_visitors.cc b/src/relay/transforms/device_aware_visitors.cc
index f3ca1bfa3..7fba7c947 100644
--- a/src/relay/transforms/device_aware_visitors.cc
+++ b/src/relay/transforms/device_aware_visitors.cc
@@ -210,6 +210,10 @@ void DeviceAwareExprVisitor::PostVisitLetBlock_(const LetNode* let_node) {
 }
 
 Expr DeviceAwareExprMutator::VisitExpr_(const FunctionNode* function_node) {
+  if (function_node->HasNonzeroAttr(attr::kObf)) {
+    // zz: no tracking inside obf function, already handled
+    return DeviceAwareVisitExpr_(function_node);
+  }
   if (function_node->HasNonzeroAttr(attr::kPrimitive)) {
     // No tracking inside primitive functions.
     return DeviceAwareVisitExpr_(function_node);
diff --git a/src/relay/transforms/fuse_ops.cc b/src/relay/transforms/fuse_ops.cc
index ee005aa17..c12620fa2 100644
--- a/src/relay/transforms/fuse_ops.cc
+++ b/src/relay/transforms/fuse_ops.cc
@@ -89,6 +89,236 @@ static const Op& stop_fusion_op = Op::Get("annotation.stop_fusion");
 TVM_REGISTER_PASS_CONFIG_OPTION("relay.FuseOps.max_depth", Integer);
 TVM_REGISTER_PASS_CONFIG_OPTION("relay.FuseOps.link_params", Bool);
 
+// zz: helper functions
+std::vector<int> get_vec_shape_from_tvm_shape(const Array<PrimExpr>& shape) {
+  std::vector<int> vec_shape;
+  for (IndexExpr dim : shape) {
+    const int64_t* pval = tir::as_const_int(dim);
+    ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << shape;
+    ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
+    vec_shape.push_back(static_cast<int>(pval[0]));
+  }
+  return vec_shape;
+}
+std::vector<int> get_divisors(int n) {
+  // def get_divisors(n):
+  //   divisors = []
+  //   for i in range(1, int(n**0.5) + 1):
+  //       if n % i == 0:
+  //           divisors.append(i)
+  //           if i != n // i:
+  //               divisors.append(n // i)
+  //   return sorted(divisors)
+
+  std::vector<int> divisors;
+  for (int i = 1; i <= static_cast<int>(std::sqrt(n)); i++) {
+    if (n % i == 0) {
+      divisors.push_back(i);
+      if (i != n / i) {
+        divisors.push_back(n / i);
+      }
+    }
+  }
+  std::sort(divisors.begin(), divisors.end());
+  return divisors;
+}
+std::vector<int> get_rand_shape(int s, int n) {
+  // def get_rand_shape(s, n):
+  //   res = []
+  //   while n != 0:
+  //       if s == 1:
+  //           res.append(1)
+  //           n = n - 1
+  //           continue
+  //       divisors = get_divisors(s)
+  //       non_trivial_divisors = []
+  //       for d in divisors:
+  //           if d != 1 and d != s:
+  //               non_trivial_divisors.append(d)
+  //       if len(non_trivial_divisors) != 0:
+  //           rand_div = random.choice(non_trivial_divisors)
+  //           res.append(rand_div)
+  //           s = s // rand_div
+  //           n = n - 1
+  //       else:
+  //           s = math.isqrt(s) ** 2
+  //   return res
+
+  std::vector<int> res;
+  while (n > 0) {
+    if (s == 1) {
+      res.push_back(1);
+      n--;
+      continue;
+    }
+    auto divisors = get_divisors(s);
+    std::vector<int> non_trivial_divisors;
+    for (int d : divisors) {
+      if (d != 1 && d != s) {
+        non_trivial_divisors.push_back(d);
+      }
+    }
+    if (!non_trivial_divisors.empty()) {
+      int rand_div = non_trivial_divisors[rand() % non_trivial_divisors.size()];
+      res.push_back(rand_div);
+      s /= rand_div;
+      n--;
+    } else {
+      s = static_cast<int>(std::sqrt(s)) * static_cast<int>(std::sqrt(s));
+    }
+  }
+  return res;
+}
+std::vector<std::vector<int>> rand_conv2d_attr(int si, int so) {
+  // def get_rand_conv2d_attr(si, so):
+  //   in1, in2, in3, in4 = (0,0,0,0)
+  //   w1, w2, w3, w4 = (0,0,0,0)
+  //   o1, o2, o3, o4 = (0,0,0,0)
+  //   kcand = [(2,2), (3,3), (5,5), (7,7)]
+  //   success = False
+  //   while not success:
+  //       w3 , w4 = random.choice(kcand)
+  //       in1, in2, in3, in4 = get_rand_shape(si, 4)
+  //       if in3 <  w3 or in4 < w4:
+  //           continue
+  //       o3 = in3 - w3 + 1
+  //       o4 = in4 - w4 + 1
+  //       o1 = in1
+  //       if o1 * o3 * o4 >= so:
+  //           continue
+  //       if random.random() < 0.5:
+  //           o2 = math.ceil(so / (o1 * o3 * o4))
+  //       else:
+  //           o2 = math.floor(so / (o1 * o3 * o4))
+
+  //       w1 = o2
+  //       w2 = in2
+  //       success = True
+  //   return [(in1, in2, in3, in4), (w1, w2, w3, w4), (o1, o2, o3, o4)]
+
+  int in1, in2, in3, in4 = 0;
+  int w1, w2, w3, w4 = 0;
+  int o1, o2, o3, o4 = 0;
+  std::vector<std::pair<int, int>> kcand = {{2, 2}, {3, 3}, {5, 5}, {7, 7}};
+
+  bool success = false;
+  int times = 0;
+  while (!success) {
+    if (times++ > 1000) {
+      std::cerr << "Warning: rand_conv2d_attr failed to find a valid configuration after 1000 attempts." << std::endl;
+      return {{0, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}};
+    }
+
+    auto rand_pair = kcand[rand() % kcand.size()];
+    w3 = rand_pair.first;
+    w4 = rand_pair.second;
+
+    auto rand_shape = get_rand_shape(si, 4);
+    in1 = rand_shape[0];
+    in2 = rand_shape[1];
+    in3 = rand_shape[2];
+    in4 = rand_shape[3];
+
+    if (in3 < w3 || in4 < w4) {
+      continue;
+    }
+    o3 = in3 - w3 + 1;
+    o4 = in4 - w4 + 1;
+    o1 = in1;
+    if (o1 * o3 * o4 >= so) {
+      continue;
+    }
+
+    if (static_cast<float>(rand()) / RAND_MAX < 0.5) {
+      o2 = static_cast<int>(std::ceil(static_cast<float>(so) / (o1 * o3 * o4)));
+    } else {
+      o2 = static_cast<int>(std::floor(static_cast<float>(so) / (o1 * o3 * o4)));
+    }
+
+    w1 = o2;
+    w2 = in2;
+    success = true;
+  }
+
+  return {{in1, in2, in3, in4}, {w1, w2, w3, w4}, {o1, o2, o3, o4}};
+}
+
+
+// zz: mutator for fake ops
+// annotate "obf_input" (single input with range info)
+class ObfCallAnnotator : public ExprMutator {
+ public:
+  // ObfCallAnnotator(IndexedForwardGraph graph) : data_dep_graph_(graph) {}
+  ObfCallAnnotator() {}
+  Expr VisitExpr_(const CallNode* call_node) final {
+    // return ExprMutator::VisitExpr_(call_node);
+    auto new_op = this->Mutate(call_node->op);
+
+    tvm::Array<Type> ty_args;
+    ty_args.reserve(call_node->type_args.size());
+
+    // for (auto ty_arg : call_node->type_args) {
+    //   auto new_ty_arg = this->VisitType(ty_arg);
+    //   ty_args.push_back(new_ty_arg);
+    // }
+
+    tvm::Array<Expr> call_args;
+    call_args.reserve(call_node->args.size());
+    Array<Expr> obf_inputs;
+    for (auto arg : call_node->args) {
+      auto new_arg = this->Mutate(arg);
+      call_args.push_back(new_arg);
+      ty_args.push_back(new_arg->checked_type());
+      if (new_arg->IsInstance<CallNode>()) {
+        const CallNode* arg_call = new_arg.as<CallNode>();
+        // if (arg_call->op.as<OpNode>()) {
+        if (arg_call->obf_attrs.defined() && arg_call->obf_attrs.HasIntAttr("obf_op_id")) {
+          // get the shape of arg_call
+          // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg_call: " << PrettyPrint(arg_call) << std::endl;
+          // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg_call->checked_type: " << arg_call->checked_type() << std::endl;
+          // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg_call->checked_type->shape: " << arg_call->checked_type().as<TensorTypeNode>()->shape << std::endl;
+          // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg_call->checked_type->shape[0]: " << arg_call->checked_type().as<TensorTypeNode>()->shape[0] << std::endl;
+          // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg_call->checked_type->shape[1]: " << arg_call->checked_type().as<TensorTypeNode>()->shape[1] << std::endl;
+          // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg_call->checked_type->shape[2]: " << arg_call->checked_type().as<TensorTypeNode>()->shape[2] << std::endl;
+          // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg_call->checked_type->shape[3]: " << arg_call->checked_type().as<TensorTypeNode>()->shape[3] << std::endl;
+          obf_inputs.push_back(new_arg);
+        }
+      }
+      else if (new_arg->IsInstance<VarNode>()) { // zz: VarNode is free var at this stage
+        // std::cout << "    (zz) ObfCallAnnotator::VisitExpr_ arg free var: " << PrettyPrint(arg) << std::endl;
+        // check VarNode is float32 type
+        const VarNode* var_node = new_arg.as<VarNode>();
+        auto var_type = var_node->type_annotation;
+        if (var_type.as<TensorTypeNode>()->dtype == DataType::Float(32)) {
+          obf_inputs.push_back(new_arg);
+        }
+      }
+    }
+    // std::cout << "[ZZDEBUG] ObfCallAnnotator::VisitExpr_ call_args.size = " << call_args.size() << ", ty_args.size = " << ty_args.size() << std::endl;
+    // std::cout << "[ZZDEBUG] ObfCallAnnotator::VisitExpr_ call_node: " << std::endl << PrettyPrint(GetRef<Expr>(call_node)) << std::endl << "[ZZDEBUG] ObfCallAnnotator::VisitExpr_ call_node end" << std::endl;
+    // std::cout << "[ZZDEBUG] ObfCallAnnotator::VisitExpr_ obf_inputs.size() = " << obf_inputs.size() << std::endl;
+    DictAttrs new_obf_attrs = call_node->obf_attrs;
+    // if (obf_inputs.size() == 1) {
+    //   // set obf_input as an obf_attr
+    //   Expr obf_input = obf_inputs[0];
+    //   std::cout << "[ZZDEBUG] obf_input: " << std::endl << PrettyPrint(obf_input) << std::endl << "[ZZDEBUG] obf_input end" << std::endl;
+    //   new_obf_attrs.CopyOnWrite()->dict.Set("obf_input", obf_input); // annotate obf_input that should have value range
+    // }
+    if (new_obf_attrs.defined()) { // zz: only set for op that can generate obf outputs
+      new_obf_attrs.CopyOnWrite()->dict.Set("obf_inputs", obf_inputs);
+    }
+    // new_obf_attrs.CopyOnWrite()->dict.Set("obf_inputs", obf_inputs);
+
+    return WithFields(GetRef<Call>(call_node), new_op, call_args, {}, ty_args, {}, {}, new_obf_attrs);
+    // Expr res = WithFields(GetRef<Call>(call_node), new_op, call_args, {}, ty_args, {}, {}, new_obf_attrs);
+    // std::cout << "[ZZDEBUG] ObfCallAnnotator::VisitExpr_ res: " << std::endl << PrettyPrint(res) << std::endl << "[ZZDEBUG] ObfCallAnnotator::VisitExpr_ res end" << std::endl;
+    // return res;
+  }
+
+//  private:
+  // IndexedForwardGraph data_dep_graph_;
+};
+
 // Creator of post dominator tree of the dataflow
 class IndexedForwardGraphCreator : private ExprVisitor {
  public:
@@ -171,6 +401,7 @@ class IndexedForwardGraphCreator : private ExprVisitor {
       // as opaque, we will not choose to fuse it.
       node->pattern = kOpaque;
     }
+    // std::cout << std::endl << "^^^^^^ IndexedForwardGraph::VisitExpr_ ConstantNode End ^^^^^^" << std::endl;
   }
 
   void VisitExpr_(const CallNode* call) final {
@@ -319,31 +550,150 @@ class IndexedForwardGraphCreator : private ExprVisitor {
 
 class FuseMutator : private MixedModeMutator {
  public:
-  FuseMutator(int fuse_opt_level, size_t max_fuse_depth, size_t max_function_args, bool link_params)
+  FuseMutator(int fuse_opt_level, size_t max_fuse_depth, size_t max_function_args, bool link_params, bool enable_flex_fuse, bool enable_fakeop_insert) // zz: add obf switch
       : fuse_opt_level_(fuse_opt_level),
         max_fuse_depth_(max_fuse_depth),
         max_function_args_(max_function_args),
-        link_params_(link_params) {}
+        // link_params_(link_params) {}
+        link_params_(link_params),
+        enable_flex_fuse_(enable_flex_fuse), // zz: add obf switch
+        enable_fakeop_insert_(enable_fakeop_insert) {} // zz: add obf switch
 
   // Run the transform
   Expr Transform(const Expr& body) {
     return Transform(body, fuse_opt_level_, max_fuse_depth_, link_params_);
   }
 
+  // zz: for random fuse, ref Transform(const Expr& body)
+  Expr TransformObf(const Expr& body) {
+    // return TransformObf(body, fuse_opt_level_, max_fuse_depth_, link_params_);
+    Expr res = TransformObf(body, fuse_opt_level_, max_fuse_depth_, link_params_);
+    // std::cout << "[ZZLOG] obf_op_ids_for_fops_: ";
+    // for (int obf_op_id : obf_op_ids_for_fops_) {
+    //   std::cout << obf_op_id << " ";
+    // }
+    // std::cout << std::endl;
+    return res;
+  }
+
  protected:
   // Run the transform
   Expr Transform(const Expr& body, int fuse_opt_level, size_t max_fuse_depth, bool link_params) {
     // setup the group map.
     auto graph = IndexedForwardGraphCreator::Create(&arena_, body);
+    // std::cout << "[ZZDEBUG] GraphPartitioner: fuse_opt_level = " << fuse_opt_level << ", max_fuse_depth = " << max_fuse_depth << ", max_function_args = " << max_function_args_ << std::endl;
     auto groups = GraphPartitioner(&arena_, fuse_opt_level, max_fuse_depth, max_function_args_)
                       .Partition(graph);
     for (size_t nid = 0; nid < graph.post_dfs_order.size(); ++nid) {
       ICHECK(graph.post_dfs_order[nid]->ref != nullptr);
       gmap_[graph.post_dfs_order[nid]->ref] = groups[nid];
     }
+    // std::cout << "[ZZDEBUG] FuseMutator Transform DebugDumpGroupStr: " << std::endl << this->DebugDumpGroupStr(body) << "[ZZDEBUG] FuseMutator Transform DebugDumpGroupStr end" << std::endl;
     // The following line can be used for debug.
     // this->DebugDumpGroup(body);
     return this->Mutate(body);
+    // std::cout << "[ZZDEBUG] Transform FuseMutator before mutate:" << std::endl << PrettyPrint(body) << std::endl;
+    // auto res = this->Mutate(body);
+    // std::cout << "[ZZDEBUG] Transform FuseMutator after mutate:" << std::endl << PrettyPrint(res) << std::endl;
+    // return res;
+  }
+
+  // zz: for random fuse, ref Transform(const Expr& body, int fuse_opt_level, size_t max_fuse_depth, bool link_params)
+  Expr TransformObf(const Expr& body, int fuse_opt_level, size_t max_fuse_depth, bool link_params) {
+    // std::cout << "[ZZDEBUG] FuseMutator at Transform:" << std::endl << PrettyPrint(body) << "[ZZDEBUG] FuseMutator at Transform end" << std::endl;
+    // setup the group map.
+    ObfCallAnnotator obf_call_annotator;
+    auto annot_body = obf_call_annotator.Mutate(body);
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator annot_body:" << std::endl << PrettyPrint(annot_body) << std::endl << "[ZZDEBUG] FuseMutator TransformObf annot_body end" << std::endl;
+    auto graph = IndexedForwardGraphCreator::Create(&arena_, annot_body);
+    // zz: print graph info
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator graph size = " << graph.post_dfs_order.size() << std::endl;
+    for (size_t i = 0; i < graph.post_dfs_order.size(); i++) {
+      auto* node = graph.post_dfs_order[i];
+      size_t ouput_size = 0;
+      for (auto* link = node->outputs.head; link != nullptr; link = link->next) {
+        ouput_size++;
+      }
+      // std::cout << "  (zz) node[" << i << "]: outputs.size()= (0 or 1 in mnist) " << ouput_size << std::endl;
+    }
+
+    /* =========================================================================================================
+       zz: split GraphPartitioner & Partition for mannual fuse testing
+       ========================================================================================================= */
+    // zz: original GraphPartitioner.Partition, comment out
+    // std::cout << std::endl << "[ZZDEBUG] FuseMutator Transform graph node_map size before = " << graph.node_map.size() << std::endl;
+    // std::cout << std::endl << "[ZZDEBUG] FuseMutator Transform graph post_dfs_order size before = " << graph.post_dfs_order.size() << std::endl;
+    // if (!skip_fuse_) {
+    // auto groups = GraphPartitioner(&arena_, fuse_opt_level, max_fuse_depth, max_function_args_)
+    //                   .Partition(graph);
+    // std::cout << std::endl << "[ZZDEBUG] FuseMutator Transform graph size after Partition = " << graph.post_dfs_order.size() << std::endl;
+    // std::cout << std::endl << "[ZZDEBUG] FuseMutator Transform group size after Partition = " << groups.size() << std::endl;
+
+    // zz: split GraphPartitioner & Partition
+    auto gp = GraphPartitioner(&arena_, fuse_opt_level, max_fuse_depth, max_function_args_);
+    gp.PartitionWithoutMov(graph);
+    // gp.MannualFuse(graph, 3); // layout_transform
+    // gp.MannualFuse(graph, 8); // fuse maxpool, make complicated fusion (conv + maxpool)
+    // gp.MannualFuse(graph, 11); // nn.pad
+    // gp.MannualFuse(graph, 19); // (layout_transform) reshape
+
+    // gp.MannualFuse(graph, 4); // naive nn model (fuse to downstream maxpool)
+
+    // gp.MannualFuse(graph, 7); // naive expr_with_if if
+
+    auto groups = gp.GetGroups();
+    auto fops = gp.GetFops();
+    /* =========================================================================================================
+       zz: end
+       ========================================================================================================= */
+
+    for (size_t nid = 0; nid < graph.post_dfs_order.size(); ++nid) {
+      ICHECK(graph.post_dfs_order[nid]->ref != nullptr);
+      gmap_[graph.post_dfs_order[nid]->ref] = groups[nid];
+      // zz: handle should_flat_nodes_
+      bool should_flat = false;
+      auto graph_node = graph.post_dfs_order[nid];
+      for (auto link = graph_node->outputs.head; link != nullptr; link = link->next) {
+        if (fops.count(link->value.node->ref)) {
+          should_flat = true;
+          break;
+        }
+      }
+      if (should_flat) {
+        should_flat_nodes_.insert(graph_node->ref);
+      }
+    }
+    // The following line can be used for debug.
+    // this->DebugDumpGroup(body);
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator before mutate:" << std::endl << PrettyPrint(annot_body) << std::endl;
+    // annot_body is function, get the body of it
+    if (annot_body->IsInstance<FunctionNode>()) {
+      last_op_ = Downcast<Function>(annot_body)->body;
+    }
+    selected_fops_ = fops; // zz: save selected fake ops from partition result
+    // print fops info
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator fops size = " << fops.size() << std::endl;
+    for (auto* fop : selected_fops_) {
+      // cast to Expr
+      Expr fop_expr = GetRef<Expr>(static_cast<const ExprNode*>(fop));
+      // std::cout << "  (zz) fop: " << std::endl << PrettyPrint(fop_expr) << std::endl;
+    }
+    // print should_flat_nodes_ info
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator should_flat_nodes_ size = " << should_flat_nodes_.size() << std::endl;
+    for (auto* node : should_flat_nodes_) {
+      // cast to Expr
+      Expr node_expr = GetRef<Expr>(static_cast<const ExprNode*>(node));
+      // std::cout << "  (zz) should_flat_node: " << std::endl << PrettyPrint(node_expr) << std::endl;
+    }
+
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator last_op_: " << std::endl << PrettyPrint(last_op_) << std::endl;
+    auto res = this->Mutate(annot_body);
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator after mutate:" << std::endl << PrettyPrint(res) << std::endl;
+    // infer type for res
+    auto res_with_type = InferType(res);
+    // std::cout << "[ZZDEBUG] TransformObf FuseMutator after mutate res_with_type:" << std::endl << PrettyPrint(res_with_type) << std::endl;
+    return res_with_type;
+    // return this->Mutate(annot_body);
   }
 
  private:
@@ -352,6 +702,11 @@ class FuseMutator : private MixedModeMutator {
   size_t max_function_args_;
   bool link_params_;
 
+  // zz: for obfuscation
+  Expr last_op_;
+  bool enable_flex_fuse_;
+  bool enable_fakeop_insert_;
+
   using MixedModeMutator::VisitExpr_;
 
   /*! \brief Temporary information from each group. */
@@ -361,6 +716,8 @@ class FuseMutator : private MixedModeMutator {
     Array<Var> params;
     // The arguments to call the functions.
     Array<Expr> arguments;
+    // zz: mark whether the group have fake ops
+    bool has_fake_ops = false;
     // Get a new parameter or allocate an old one
     Var GetOrAllocParam(const Expr& expr, const Type& type) {
       // run linear scan as most fused groups contain only a few inputs.
@@ -382,18 +739,654 @@ class FuseMutator : private MixedModeMutator {
   std::unordered_map<const Object*, GraphPartitioner::Group*> gmap_;
   /* \brief Internal group information map. */
   std::unordered_map<GraphPartitioner::Group*, GroupInfo> ginfo_;
+  // zz: selected fake ops
+  std::unordered_set<const Object*> selected_fops_;
+  // zz: nodes should be flat because output is fop
+  std::unordered_set<const Object*> should_flat_nodes_;
+  // zz: for logging obf op id that is used to construct fop
+  std::unordered_set<int> obf_op_ids_for_fops_;
+
+  // zz: needed for model input handling
+  Expr VisitExpr_(const VarNode* var_node) {
+    if (!enable_fakeop_insert_) {
+      return ExprMutator::VisitExpr_(var_node);
+    }
+    // reshape model input to 1D flatten tensor
+    // ref: MakeReshape
+    auto flat_op = Op::Get("reshape");
+    auto attrs = make_object<ReshapeAttrs>();
+    // get size of var_node
+    auto shape = var_node->checked_type().as<TensorTypeNode>()->shape;
+    int flatten_size = 1;
+    for (IndexExpr dim : shape) {
+      const int64_t* pval = tir::as_const_int(dim);
+      ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << shape;
+      ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
+      flatten_size *= static_cast<size_t>(pval[0]);
+    }
+    auto newshape = Array<Integer>({Integer(flatten_size)});
+    attrs->newshape = std::move(newshape);
+
+
+    auto obf_attrs_node = make_object<DictAttrsNode>();
+    obf_attrs_node->dict.Set("obf_flatten_output", tvm::Integer(1));
+    auto new_params = Array<Var>({Var("input", var_node->checked_type())});
+    // auto call = Call(flat_op, {GetRef<Var>(var_node)}, Attrs(), {}, {}, DictAttrs(obfa_attrs_node));
+    auto call = Call(flat_op, {new_params[0]}, Attrs(attrs), {}, {}, DictAttrs(obf_attrs_node));
+    // infer type for call
+    // auto call_with_type = InferType(call);
+    // set call checked_type
+    Type call_type = TensorType({flatten_size}, var_node->checked_type().as<TensorTypeNode>()->dtype);
+    call->checked_type_ = call_type;
+    // std::cout << "[ZZDEBUG] FuseMutator VisitExpr_ VarNode final call:" << std::endl << PrettyPrint(call) << std::endl << "[ZZDEBUG] FuseMutator VisitExpr_ VarNode final call end" << std::endl;
+    // std::cout << "[ZZDEBUG] FuseMutator VisitExpr_ VarNode final call_with_type:" << std::endl << PrettyPrint(call_with_type) << std::endl << "[ZZDEBUG] FuseMutator VisitExpr_ VarNode final call_with_type end" << std::endl;
+    // return call_with_type;
+    // make a new function for call_with_type
+    auto new_body = call;
+    auto new_ret_type = call->checked_type();
+    auto new_ty_params = Array<TypeVar>();
+    auto new_func = Function(new_params, new_body, new_ret_type, new_ty_params, {});
+    new_func = WithAttr(std::move(new_func), attr::kPrimitive, tvm::Integer(1));
+    // call new_func
+    auto new_call = Call(new_func, {GetRef<Var>(var_node)}, Attrs(), {}, {});
+    // infer type for new_call
+    auto new_call_with_type = InferType(new_call);
+    return new_call_with_type;
+  }
 
   // Skip primitive function.
   Expr VisitExpr_(const FunctionNode* fn_node) {
     if (fn_node->HasNonzeroAttr(attr::kPrimitive)) {
       return GetRef<Expr>(fn_node);
     } else {
-      return ExprMutator::VisitExpr_(fn_node);
+      // return ExprMutator::VisitExpr_(fn_node);
+      if (!enable_fakeop_insert_) {
+        return ExprMutator::VisitExpr_(fn_node);
+      }
+      // zz: Var Node can become call node due to reshape above
+      tvm::Array<TypeVar> ty_params;
+
+      for (auto ty_param : fn_node->type_params) {
+        TypeVar new_ty_param = Downcast<TypeVar>(VisitType(ty_param));
+        ty_params.push_back(new_ty_param);
+      }
+
+      tvm::Array<Var> params;
+      for (auto param : fn_node->params) {
+        // Var new_param = Downcast<Var>(this->Mutate(param));
+        Expr new_param = this->Mutate(param);
+        if (new_param->IsInstance<VarNode>()) {
+          params.push_back(Downcast<Var>(new_param));
+        }
+        else {
+          params.push_back(param);
+        }
+
+      }
+
+      auto ret_type = this->VisitType(fn_node->ret_type);
+      auto body = this->Mutate(fn_node->body);
+
+      return WithFields(GetRef<Function>(fn_node), params, body, ret_type, ty_params);
+    }
+  }
+
+  // zz: obf helper function, legalize the input so that it can be used as an argument to call
+  // basically adding slice & reshape
+  Expr GetNewObfInput(const CallNode* call, const Expr obf_input, GraphPartitioner::Group* ret_group) {
+      // check obf_input should be one of the call args
+      bool found = false;
+      for (auto arg : call->args) {
+        if (obf_input.same_as(arg)) {
+          found = true;
+          break;
+        }
+      }
+      ICHECK(found) << "obf_input should be one of the call args";
+
+      // std::cout << "[ZZDEBUG] FuseMutator GetNewObfInput obf_input: " << std::endl << PrettyPrint(obf_input) << std::endl;
+      Expr new_obf_input = this->VisitExpr(obf_input);
+      // std::cout << "[ZZDEBUG] FuseMutator GetNewObfInput new_obf_input: " << std::endl << PrettyPrint(new_obf_input) << std::endl;
+      // get new_obf_input shape
+      auto ori_shape = obf_input->checked_type().as<TensorTypeNode>()->shape;
+      auto ori_shape_vec = get_vec_shape_from_tvm_shape(ori_shape);
+      int ori_size = 1;
+      for (size_t i = 0; i < ori_shape_vec.size(); i++) {
+        ori_size *= ori_shape_vec[i];
+      }
+      auto new_shape = new_obf_input->checked_type().as<TensorTypeNode>()->shape;
+      auto new_shape_vec = get_vec_shape_from_tvm_shape(new_shape);
+      int new_size = 1;
+      for (size_t i = 0; i < new_shape_vec.size(); i++) {
+        new_size *= new_shape_vec[i];
+      }
+      // check if we need reshape here
+      // ICHECK(new_shape_vec.size() == 1); // flatten tensor (zz: input not necessary to be flatten anymore)
+      ICHECK(ori_size <= new_size);
+      bool need_slice = false;
+      bool need_reshape = false;
+      if (ori_size < new_size) {
+        need_reshape = true;
+        need_slice = true;
+      }
+      if (!need_reshape) {
+        if (ori_shape_vec.size() != new_shape_vec.size()) {
+          need_reshape = true;
+        }
+        else {
+          for (size_t i = 0; i < ori_shape_vec.size(); i++) {
+            if (ori_shape_vec[i] != new_shape_vec[i]) {
+              need_reshape = true;
+              break;
+            }
+          }
+        }
+      }
+
+      Expr obf_input_before_slice = new_obf_input;
+      Expr obf_call_arg = NullValue<Expr>();
+      // insert slice and reshape
+      if (need_slice) {
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode need_slice" << std::endl;
+        // create slice op
+        auto slice_op = Op::Get("strided_slice");
+        auto slice_attrs = make_object<StridedSliceAttrs>();
+        // set slice_attrs
+        slice_attrs->begin = Array<Integer>({Integer(0)});
+        slice_attrs->end = Array<Integer>({Integer(ori_size)});
+        slice_attrs->strides = Array<Integer>({Integer(1)});
+        // create slice call
+        obf_call_arg = GetNewArgument(obf_input, ret_group);
+        Expr slice_obf_call_arg = obf_call_arg;
+        auto slice_call = Call(slice_op, {slice_obf_call_arg}, Attrs(slice_attrs), {}, {}, DictAttrs());
+
+        // // set slice_call checked_type
+        // Type slice_call_type = TensorType({ori_size}, obf_input->checked_type().as<TensorTypeNode>()->dtype);
+        // slice_call->checked_type_ = slice_call_type;
+
+        // infer type for slice_call
+        auto slice_call_with_type = InferType(slice_call);
+        slice_call->checked_type_ = slice_call_with_type->checked_type();
+
+        new_obf_input = slice_call;
+      }
+      if (need_reshape) {
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode need_reshape" << std::endl;
+        // create reshape op
+        auto reshape_op = Op::Get("reshape");
+        auto reshape_attrs = make_object<ReshapeAttrs>();
+        // set reshape_attrs
+        auto newshape = Array<Integer>();
+        for (size_t i = 0; i < ori_shape_vec.size(); i++) {
+          newshape.push_back(Integer(ori_shape_vec[i]));
+        }
+        reshape_attrs->newshape = std::move(newshape);
+        // create reshape call
+        Expr reshape_obf_call_arg = new_obf_input;
+        if (new_obf_input == obf_input_before_slice) {
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode reshape new_obf_input == obf_input" << std::endl;
+          obf_call_arg = GetNewArgument(obf_input, ret_group);
+          reshape_obf_call_arg = obf_call_arg;
+        }
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode reshape reshape_obf_call_arg:" << std::endl << PrettyPrint(reshape_obf_call_arg) << std::endl << "[ZZDEBUG] FuseMutator Rewrite_ CallNode reshape reshape_obf_call_arg end" << std::endl;
+        auto reshape_call = Call(reshape_op, {reshape_obf_call_arg}, Attrs(reshape_attrs), {}, {}, DictAttrs());
+
+        // // set reshape_call checked_type
+        // Type reshape_call_type = TensorType(ori_shape, obf_input->checked_type().as<TensorTypeNode>()->dtype);
+        // reshape_call->checked_type_ = reshape_call_type;
+
+        // infer type for reshape_call
+        auto reshape_call_with_type = InferType(reshape_call);
+        reshape_call->checked_type_ = reshape_call_with_type->checked_type();
+
+        new_obf_input = reshape_call;
+      }
+      // ICHECK(obf_call_arg.defined()); // zz: not necessary anymore
+
+      return new_obf_input;
+  }
+  // [TODO] zz: ref GetNewObfInput, merge code if possible
+  Expr GetNewObfInputForFakeOp(std::vector<int> in_shape, const Expr obf_input, GraphPartitioner::Group* ret_group) {
+    Expr new_obf_input = this->VisitExpr(obf_input);
+    auto ori_shape_vec = in_shape;
+    int ori_size = 1;
+    for (size_t i = 0; i < ori_shape_vec.size(); i++) {
+      ori_size *= ori_shape_vec[i];
+    }
+    auto new_shape = new_obf_input->checked_type().as<TensorTypeNode>()->shape;
+    auto new_shape_vec = get_vec_shape_from_tvm_shape(new_shape);
+    int new_size = 1;
+    for (size_t i = 0; i < new_shape_vec.size(); i++) {
+      new_size *= new_shape_vec[i];
+    }
+    // check if we need reshape here
+    // ICHECK(new_shape_vec.size() == 1); // flatten tensor (zz: input not necessary to be flatten anymore)
+    ICHECK(ori_size <= new_size);
+    bool need_slice = false;
+    bool need_reshape = false;
+    if (ori_size < new_size) {
+      need_reshape = true;
+      need_slice = true;
+    }
+    if (!need_reshape) {
+      if (ori_shape_vec.size() != new_shape_vec.size()) {
+        need_reshape = true;
+      }
+      else {
+        for (size_t i = 0; i < ori_shape_vec.size(); i++) {
+          if (ori_shape_vec[i] != new_shape_vec[i]) {
+            need_reshape = true;
+            break;
+          }
+        }
+      }
+    }
+
+    Expr obf_input_before_slice = new_obf_input;
+    Expr obf_call_arg = NullValue<Expr>();
+    // insert slice and reshape
+    if (need_slice) {
+      // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode need_slice" << std::endl;
+      // create slice op
+      auto slice_op = Op::Get("strided_slice");
+      auto slice_attrs = make_object<StridedSliceAttrs>();
+      // set slice_attrs
+      slice_attrs->begin = Array<Integer>({Integer(0)});
+      slice_attrs->end = Array<Integer>({Integer(ori_size)});
+      slice_attrs->strides = Array<Integer>({Integer(1)});
+      // create slice call
+      obf_call_arg = GetNewArgument(obf_input, ret_group);
+      Expr slice_obf_call_arg = obf_call_arg;
+      auto slice_call = Call(slice_op, {slice_obf_call_arg}, Attrs(slice_attrs), {}, {}, DictAttrs());
+
+      // // set slice_call checked_type
+      // Type slice_call_type = TensorType({ori_size}, obf_input->checked_type().as<TensorTypeNode>()->dtype);
+      // slice_call->checked_type_ = slice_call_type;
+
+      // infer type for slice_call
+      auto slice_call_with_type = InferType(slice_call);
+      slice_call->checked_type_ = slice_call_with_type->checked_type();
+
+      new_obf_input = slice_call;
+    }
+    if (need_reshape) {
+      // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode need_reshape" << std::endl;
+      // create reshape op
+      auto reshape_op = Op::Get("reshape");
+      auto reshape_attrs = make_object<ReshapeAttrs>();
+      // set reshape_attrs
+      auto newshape = Array<Integer>();
+      for (size_t i = 0; i < ori_shape_vec.size(); i++) {
+        newshape.push_back(Integer(ori_shape_vec[i]));
+      }
+      reshape_attrs->newshape = std::move(newshape);
+      // create reshape call
+      Expr reshape_obf_call_arg = new_obf_input;
+      if (new_obf_input == obf_input_before_slice) {
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode reshape new_obf_input == obf_input" << std::endl;
+        obf_call_arg = GetNewArgument(obf_input, ret_group);
+        reshape_obf_call_arg = obf_call_arg;
+      }
+      // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode reshape reshape_obf_call_arg:" << std::endl << PrettyPrint(reshape_obf_call_arg) << std::endl << "[ZZDEBUG] FuseMutator Rewrite_ CallNode reshape reshape_obf_call_arg end" << std::endl;
+      auto reshape_call = Call(reshape_op, {reshape_obf_call_arg}, Attrs(reshape_attrs), {}, {}, DictAttrs());
+
+      // // set reshape_call checked_type
+      // Type reshape_call_type = TensorType(ori_shape, obf_input->checked_type().as<TensorTypeNode>()->dtype);
+      // reshape_call->checked_type_ = reshape_call_type;
+
+      // infer type for reshape_call
+      auto reshape_call_with_type = InferType(reshape_call);
+      reshape_call->checked_type_ = reshape_call_with_type->checked_type();
+
+      new_obf_input = reshape_call;
+    }
+    // ICHECK(obf_call_arg.defined()); // zz: not necessary anymore
+
+    return new_obf_input;
+  }
+
+  Expr RandReluOp(Expr obf_input, int true_op_out_size, GraphPartitioner::Group* ret_group) {
+    Expr obf_call_arg = GetNewArgument(obf_input, ret_group);
+
+    auto fake_op = Op::Get("nn.relu");
+    auto fake_op_obf_attrs = make_object<DictAttrsNode>();
+    auto obf_call = Call(fake_op, {obf_call_arg}, Attrs(), {}, {}, {});
+
+    return obf_call;
+  }
+
+  Expr RandConv2DOp(Expr obf_input, int so, GraphPartitioner::Group* ret_group) {
+    // std::cout << "[ZZDEBUG] FuseMutator RandConv2DOp obf_input: " << std::endl << PrettyPrint(obf_input) << std::endl;
+    Expr obf_call_arg = GetNewArgument(obf_input, ret_group);
+
+    // get random attrs
+    int si = 1;
+    auto obf_call_arg_shape = NullValue<Array<IndexExpr>>();
+    if (obf_call_arg->IsInstance<VarNode>()) { // zz: var node type is not populated, use type_annotation from VarNode
+      auto var_node = obf_call_arg.as<VarNode>();
+      obf_call_arg_shape = var_node->type_annotation.as<TensorTypeNode>()->shape;
+    }
+    else {
+      obf_call_arg_shape = obf_call_arg->checked_type().as<TensorTypeNode>()->shape;
+    }
+    for (IndexExpr dim : obf_call_arg_shape) {
+      const int64_t* pval = tir::as_const_int(dim);
+      ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << obf_call_arg_shape;
+      ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
+      si *= static_cast<size_t>(pval[0]);
+    }
+
+    std::vector<std::vector<int>> conv2d_attrs = rand_conv2d_attr(si, so);
+    std::vector<int> in_shape = conv2d_attrs[0];
+    std::vector<int> w_shape = conv2d_attrs[1];
+    std::vector<int> out_shape = conv2d_attrs[2];
+    int in_size = in_shape[0] * in_shape[1] * in_shape[2] * in_shape[3];
+    int w_size = w_shape[0] * w_shape[1] * w_shape[2] * w_shape[3];
+    int out_size = out_shape[0] * out_shape[1] * out_shape[2] * out_shape[3];
+    if (in_size == 0 || w_size == 0 || out_size == 0) {
+      return NullValue<Expr>(); // fake op generation failed
+    }
+    // else {
+    //   std::cout << "[ZZDEBUG] RandConv2DOp in_shape: " << in_shape[0] << ", " << in_shape[1] << ", " << in_shape[2] << ", " << in_shape[3] << std::endl;
+    //   std::cout << "[ZZDEBUG] RandConv2DOp w_shape: " << w_shape[0] << ", " << w_shape[1] << ", " << w_shape[2] << ", " << w_shape[3] << std::endl;
+    //   std::cout << "[ZZDEBUG] RandConv2DOp out_shape: " << out_shape[0] << ", " << out_shape[1] << ", " << out_shape[2] << ", " << out_shape[3] << std::endl;
+    // }
+
+    // get input expr with in_shape
+    Expr new_obf_input = GetNewObfInputForFakeOp(in_shape, obf_input, ret_group);
+
+
+    // create conv2d attrs
+    // ref relay.op.nn._make.conv2d, MakeConv<Conv2DAttrs>
+    auto fake_op = Op::Get("nn.conv2d");
+    auto fake_op_obf_attrs = make_object<Conv2DAttrs>();
+    // strides=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1
+    fake_op_obf_attrs->strides = Array<IndexExpr>({Integer(1), Integer(1)});
+    fake_op_obf_attrs->padding = Array<IndexExpr>({Integer(0), Integer(0)});
+    fake_op_obf_attrs->dilation = Array<IndexExpr>({Integer(1), Integer(1)});
+    fake_op_obf_attrs->groups = 1;
+    fake_op_obf_attrs->channels = out_shape[1];
+    fake_op_obf_attrs->kernel_size = Array<IndexExpr>({Integer(w_shape[2]), Integer(w_shape[3])});
+    // set fake_op_obf_attrs->data_layout, weight_layout, kernel_layout
+    fake_op_obf_attrs->data_layout = "NCHW";
+    fake_op_obf_attrs->kernel_layout = "OIHW";
+    fake_op_obf_attrs->out_layout = "";
+    fake_op_obf_attrs->out_dtype = DataType::Float(32);
+
+    // create const weight
+    auto weight_shape = std::vector<int64_t>({w_shape[0], w_shape[1], w_shape[2], w_shape[3]});
+    auto weight_nd = runtime::NDArray::Empty(weight_shape, DataType::Float(32), {kDLCPU, 0});
+    // fill weight_nd with random values
+    std::vector<float> weight_data(w_size);
+    for (int i = 0; i < w_size; i++) {
+      weight_data[i] = static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
+    }
+    weight_nd.CopyFromBytes(weight_data.data(), w_size * sizeof(float));
+    // create weight constant expr
+    auto weight_expr = Constant(weight_nd);
+    // std::cout << "[ZZDEBUG] FuseMutator RandConv2DOp weight_expr:" << std::endl << PrettyPrint(weight_expr) << std::endl << "[ZZDEBUG] FuseMutator RandConv2DOp weight_expr end" << std::endl;
+
+
+    // create conv2d call
+    // auto obf_call = Call(fake_op, {obf_call_arg}, Attrs(), {}, {}, {});
+    auto obf_call = Call(fake_op, {new_obf_input, weight_expr}, Attrs(fake_op_obf_attrs), {}, {}, DictAttrs());
+    // infer type for obf_call
+    // std::cout << "[ZZDEBUG] FuseMutator RandConv2DOp obf_call:" << std::endl << PrettyPrint(obf_call) << std::endl << "[ZZDEBUG] FuseMutator RandConv2DOp obf_call end" << std::endl;
+    auto obf_call_with_type = InferType(obf_call);
+    // std::cout << "[ZZDEBUG] FuseMutator RandConv2DOp obf_call_with_type:" << std::endl << PrettyPrint(obf_call_with_type) << std::endl << "[ZZDEBUG] FuseMutator RandConv2DOp obf_call_with_type end" << std::endl;
+    // set obf_call checked_type
+    // obf_call->checked_type_ = obf_call_with_type->checked_type();
+
+    // check type with out_shape
+    // auto ir_type = obf_call->checked_type().as<TensorTypeNode>();
+    // ICHECK(ir_type != nullptr) << "obf_call should have a TensorType";
+    // auto out_shape_expr = ir_type->shape;
+    // std::vector<int64_t> out_shape_vec;
+
+
+
+    return obf_call;
+  }
+
+  // zz: obf helper function, generate fake op for obf_input
+  Expr GetNewFakeOp(Expr obf_input, int true_op_out_size, GraphPartitioner::Group* ret_group) {
+    // Expr obf_call_arg = GetNewArgument(obf_input, ret_group);
+    // auto flat_input_shape = obf_call_arg->checked_type().as<TensorTypeNode>()->shape;
+    // ICHECK_EQ(flat_input_shape.size(), 1);
+    // int flat_input_size = *tir::as_const_int(flat_input_shape[0]);
+
+    // // create fake op, currently only relu
+    // auto fake_op = Op::Get("nn.relu");
+    // auto fake_op_obf_attrs = make_object<DictAttrsNode>();
+    // // fake_op_obf_attrs->dict.Set("obf_op_input", obf_input); // for relu, no reshape is needed
+    // auto obf_call = Call(fake_op, {obf_call_arg}, Attrs(), {}, {}, {});
+    // // auto obf_call = Call(obf_op, {obf_input}, Attrs(), {}, {});
+    // // run type inference for obf_call
+    // // auto obf_call_with_type = InferType(obf_call);
+
+    Expr obf_call = RandReluOp(obf_input, true_op_out_size, ret_group);
+    // Expr obf_call = RandConv2DOp(obf_input, true_op_out_size, ret_group);
+    if (!obf_call.defined())
+      obf_call = RandReluOp(obf_input, true_op_out_size, ret_group);
+
+    // return obf_call_with_type;
+    return obf_call;
+  }
+
+  Expr GetNewCondition(Expr obf_input, float neuron_range, Array<Integer> indices, GraphPartitioner::Group* ret_group) {
+    auto ori_shape = obf_input->checked_type().as<TensorTypeNode>()->shape;
+    auto ori_shape_vec = get_vec_shape_from_tvm_shape(ori_shape);
+    ICHECK(ori_shape_vec.size() == indices.size());
+    // if (ori_shape_vec.size() != indices.size()) {
+    //   std::cout << "[BUG] FuseMutator Rewrite_ CallNode obf_input:" << std::endl << PrettyPrint(obf_input) << std::endl;
+    //   std::cout << "[BUG] FuseMutator Rewrite_ CallNode call:" << std::endl << PrettyPrint(GetRef<Call>(call)) << std::endl;
+    //   std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode ori_shape_vec.size(): " << ori_shape_vec.size() << std::endl;
+    //   std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode indices.size(): " << indices.size() << std::endl;
+    //   ICHECK(ori_shape_vec.size() == indices.size());
+    // }
+
+    // get flat index from indices
+    int flat_index = 0;
+    int stride = 1;
+    for (int i = indices.size() - 1; i >= 0; i--) {
+      int index = indices[i].IntValue();
+      int dim = ori_shape_vec[i];
+      ICHECK(index < dim);
+      flat_index += index * stride;
+      stride *= dim;
+    }
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode flat_index: " << flat_index << std::endl;
+
+    Expr obf_call_arg = GetNewArgument(obf_input, ret_group);
+    // auto flat_shape = obf_call_arg->checked_type().as<TensorTypeNode>()->shape;
+    // ICHECK(flat_shape.size() == 1);
+    // int flat_size = flat_shape[0].as<IntImmNode>()->value;
+    // ICHECK(flat_index < flat_size);
+    // get flat_index element from obf_call_arg using gather_nd, ref MakeGatherND
+    auto gather_nd_op = Op::Get("gather_nd");
+    auto gather_nd_attrs = make_object<GatherNDAttrs>();
+    gather_nd_attrs->batch_dims = 0;
+    gather_nd_attrs->index_rank = NullValue<Integer>();
+    // create runtime::NDArray for indices, ref ndarray.py: def array(arr, device=cpu(0), mem_scope=None)
+    auto indices_nd = runtime::NDArray::Empty({1}, DataType::Int(32), {kDLCPU, 0});
+    indices_nd.CopyFromBytes(&flat_index, 4);
+    // create indices constant expr
+    auto indices_expr = Constant(indices_nd);
+    // create gather_nd call
+    Expr selected_ele = Call(gather_nd_op, {obf_call_arg, indices_expr}, Attrs(gather_nd_attrs), {}, {}, DictAttrs());
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode selected_ele: " << std::endl << PrettyPrint(selected_ele) << std::endl;
+    // infer type for selected_ele (type is float32)
+    auto selected_ele_with_type = InferType(selected_ele);
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode selected_ele_with_type: " << std::endl << PrettyPrint(selected_ele_with_type) << std::endl;
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode selected_ele_with_type type: " << std::endl << PrettyPrint(selected_ele_with_type->checked_type()) << std::endl;
+    // // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode selected_ele_with_type type key: " << std::endl << selected_ele_with_type->checked_type()->GetTypeKey() << std::endl; // TensorType
+    // auto ele_type = selected_ele_with_type->checked_type().as<TensorTypeNode>();
+    // auto ele_shape = ele_type->shape;
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode selected_ele_with_type shape: " << std::endl << PrettyPrint(ele_shape) << std::endl; // shape = []
+    selected_ele->checked_type_ = selected_ele_with_type->checked_type(); // TensorType, shape = []
+
+    // create constant for neuron_range
+    auto neuron_range_nd = runtime::NDArray::Empty({}, DataType::Float(32), {kDLCPU, 0});
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode neuron_range: " << neuron_range << std::endl;
+    neuron_range_nd.CopyFromBytes(&neuron_range, 4);
+    auto neuron_range_const = Constant(neuron_range_nd);
+
+    // create relay condition comparing selected_ele with neuron_range
+    Expr condition = Call(Op::Get("less"), {selected_ele, neuron_range_const}, Attrs(), {}, {}, DictAttrs());
+    // infer type for condition
+    auto condition_with_type = InferType(condition);
+    condition->checked_type_ = condition_with_type->checked_type();
+    return condition;
+  }
+
+  Expr CreateIfExpr(Array<Expr> ops, Array<Expr> conditions, size_t index) {
+    Expr if_condition = conditions[index];
+    Expr if_true = ops[index];
+    if (index == conditions.size() - 1) {
+      Expr if_false = ops[index+1];
+      Expr if_expr = If(if_condition, if_true, if_false);
+      auto if_expr_with_type = InferType(if_expr);
+      if_expr->checked_type_ = if_expr_with_type->checked_type();
+      return if_expr;
+    }
+    else {
+      Expr if_false = CreateIfExpr(ops, conditions, index+1);
+      Expr if_expr = If(if_condition, if_true, if_false);
+      auto if_expr_with_type = InferType(if_expr);
+      if_expr->checked_type_ = if_expr_with_type->checked_type();
+      return if_expr;
+    }
+  }
+
+  Expr GetFlattenOuput(Expr ori_op, int flatten_size, GraphPartitioner::Group* ret_group) {
+    // get ori_op size
+    auto ori_out_shape = ori_op->checked_type().as<TensorTypeNode>()->shape;
+    int ori_out_size = 1;
+    for (IndexExpr dim : ori_out_shape) {
+      const int64_t* pval = tir::as_const_int(dim);
+      ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << ori_out_shape;
+      ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
+      ori_out_size *= static_cast<size_t>(pval[0]);
     }
+    bool need_reshape = false;
+    bool need_pad = false;
+    if (ori_out_shape.size() != 1) {
+      need_reshape = true;
+    }
+    if (flatten_size > ori_out_size) {
+      need_pad = true;
+    }
+
+    Expr new_op = ori_op;
+    if (need_reshape) {
+      // std::cout << "[ZZDEBUG] FuseMutator GetFlattenOuput need_reshape" << std::endl;
+      // create reshape op
+      auto reshape_op = Op::Get("reshape");
+      auto reshape_attrs = make_object<ReshapeAttrs>();
+      // set reshape_attrs
+      auto newshape = Array<Integer>({Integer(ori_out_size)});
+      reshape_attrs->newshape = std::move(newshape);
+      // create reshape call
+      auto reshape_call = Call(reshape_op, {new_op}, Attrs(reshape_attrs), {}, {}, DictAttrs());
+      // infer type for reshape_call
+      auto reshape_call_with_type = InferType(reshape_call);
+      reshape_call->checked_type_ = reshape_call_with_type->checked_type();
+      new_op = reshape_call;
+    }
+    if (need_pad) {
+      // std::cout << "[ZZDEBUG] FuseMutator GetFlattenOuput need_pad" << std::endl;
+      // create pad op, ref: MakePad
+      auto pad_op = Op::Get("nn.pad");
+      auto pad_attrs = make_object<PadAttrs>();
+      // set pad_attrs
+      auto pad_width = Array<Array<Integer>>();
+      Array<Integer> pad_width_dim = {Integer(0), Integer(flatten_size - ori_out_size)};
+      pad_width.push_back(pad_width_dim);
+      pad_attrs->pad_width = std::move(pad_width);
+      pad_attrs->pad_mode = "constant";
+
+      // create constant 0.0 pad_value
+      float pad_value = 0.0;
+      auto pad_nd = runtime::NDArray::Empty({}, DataType::Float(32), {kDLCPU, 0});
+      pad_nd.CopyFromBytes(&pad_value, 4);
+      auto pad_value_expr = Constant(pad_nd);
+      auto pad_call = Call(pad_op, {new_op, pad_value_expr}, Attrs(pad_attrs), {}, {}, DictAttrs());
+
+      // infer type for pad_call
+      auto pad_call_with_type = InferType(pad_call);
+      pad_call->checked_type_ = pad_call_with_type->checked_type();
+      new_op = pad_call;
+    }
+
+    // std::cout << "[ZZDEBUG] FuseMutator GetFlattenOuput new_op:" << std::endl << PrettyPrint(new_op) << std::endl;
+
+    return new_op;
   }
 
+
+
   // Transform calls.
   Expr Rewrite_(const CallNode* call, const Expr& post) {
+    // std::cout << "================================================================================================================================================================================" << std::endl;
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode:" << std::endl << PrettyPrint(GetRef<Expr>(call)) << std::endl;
+    if(!enable_fakeop_insert_) {
+      // zz: the Rewrite_ original code
+      if (call->op.as<OpNode>()) {
+        static auto fnoncomputational = Op::GetAttrMap<TNonComputational>("TNonComputational");
+        static auto fqnncanonicalize = Op::GetAttrMap<FTVMLegalize>("FTVMQnnCanonicalize");
+
+        Op op = Downcast<Op>(call->op);
+        if (fnoncomputational.get(op, false) && !fqnncanonicalize.count(op)) {
+          return ExprMutator::VisitExpr_(call);
+        }
+
+        // If it is a primitive op call
+        // then we must have a group assignment for it already.
+        ICHECK(gmap_.count(call));
+        if (call->op == stop_fusion_op) {
+          return ExprMutator::VisitExpr(call->args[0]);
+        }
+        auto* ret_group = gmap_.at(call)->FindRoot();
+        Array<Expr> new_args = GetNewArguments(call->args, ret_group);
+
+        auto new_call = Call(call->op, new_args, call->attrs, call->type_args, call->span);
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode new_call:" << std::endl << PrettyPrint(new_call) << std::endl;
+  
+        if (ret_group->root_ref == call) {
+          // This is the root of the group
+          // create the new call node.
+          // auto ret = MakeNewFunction(ret_group, call->checked_type(), new_call);
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode ret (MakeNewFunction):" << std::endl << PrettyPrint(ret) << std::endl;
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode original call:" << std::endl << PrettyPrint(GetRef<Expr>(call)) << std::endl;
+          return MakeNewFunction(ret_group, call->checked_type(), new_call);
+        } else {
+          // This is an intermediate node of a fused function
+          // simply return the new call.
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode ret (new_call):" << std::endl << PrettyPrint(new_call) << std::endl;
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode original call:" << std::endl << PrettyPrint(GetRef<Expr>(call)) << std::endl;
+          return std::move(new_call);
+        }
+      } else {
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ hit op function" << std::endl;
+        // auto ret = ExprMutator::VisitExpr_(call);
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode ret (op function):" << std::endl << PrettyPrint(ret) << std::endl;
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode original call:" << std::endl << PrettyPrint(GetRef<Expr>(call)) << std::endl;
+        // return ret;
+        return ExprMutator::VisitExpr_(call);
+      }
+    }
+
+    // zz: obf hanlding start
+    auto ori_out_shape = call->checked_type().as<TensorTypeNode>()->shape;
+    // std::cout << "[ZZDEBUG] ori_out_shape:" << std::endl << PrettyPrint(ori_out_shape) << std::endl;
+    int ori_out_size = 1;
+    for (IndexExpr dim : ori_out_shape) {
+      const int64_t* pval = tir::as_const_int(dim);
+      ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << ori_out_shape;
+      ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
+      ori_out_size *= static_cast<size_t>(pval[0]);
+    }
+
     if (call->op.as<OpNode>()) {
       static auto fnoncomputational = Op::GetAttrMap<TNonComputational>("TNonComputational");
       static auto fqnncanonicalize = Op::GetAttrMap<FTVMLegalize>("FTVMQnnCanonicalize");
@@ -410,25 +1403,173 @@ class FuseMutator : private MixedModeMutator {
         return ExprMutator::VisitExpr(call->args[0]);
       }
       auto* ret_group = gmap_.at(call)->FindRoot();
+
+      // =========================================================================================================
+      // [zz] step 1: update new_call (the true op) with potential flattened obf input
+      // =========================================================================================================
+      auto call_obf_attrs = call->obf_attrs;
+      Array<Expr> obf_inputs = call_obf_attrs.GetAttr<Array<Expr>>("obf_inputs").value_or(Array<Expr>());
+      Array<Expr> new_obf_inputs = Array<Expr>();
+      for (size_t i = 0; i < obf_inputs.size(); i++) {
+        auto obf_input = obf_inputs[i];
+        auto new_obf_input = GetNewObfInput(call, obf_input, ret_group);
+        new_obf_inputs.push_back(new_obf_input);
+      }
+
+      // zz: update obf_input with new_obf_input (after slice & reshape)
       Array<Expr> new_args = GetNewArguments(call->args, ret_group);
+      for (auto obf_input : obf_inputs) {
+        for (size_t i = 0; i < new_args.size(); i++) {
+          // if (obf_input.defined() && obf_call_arg.defined()) {
+          if (obf_input.defined()) {
+            Expr obf_call_arg = GetNewArgument(obf_input, ret_group);
+            if (obf_call_arg.defined() && new_args[i] == obf_call_arg) {
+              // std::cout << "  (zz) FuseMutator Rewrite_ CallNode: set new_args = obf_call_arg" << std::endl;
+              // std::cout << "    (zz) from: " << std::endl << PrettyPrint(new_args[i]) << std::endl;
+              // std::cout << "    (zz) to: " << std::endl << PrettyPrint(new_obf_inputs[i]) << std::endl;
+              new_args.Set(i, new_obf_inputs[i]);
+            }
+          }
+        }
+      }
+
+      Expr new_call = Call(call->op, new_args, call->attrs, call->type_args, call->span, call_obf_attrs);
+      auto new_call_with_type = InferType(new_call); // zz: put this to MakeNewFunction
+      new_call->checked_type_ = new_call_with_type->checked_type();
+      // std::cout << "  (zz) FuseMutator Rewrite_ CallNode new_call_with_type:" << std::endl << PrettyPrint(new_call_with_type) << std::endl << "  (zz) FuseMutator Rewrite_ CallNode new_call_with_type end" << std::endl;
+
+      // =========================================================================================================
+      // [zz] step 2: generate fake ops
+      // =========================================================================================================
+      Expr obf_input = NullValue<Expr>();
+      if (obf_inputs.size() > 0) {
+        obf_input = obf_inputs[0]; // zz: only use one input to generate fake ops (only add in examples, should removed from graph partition)
+      }
+
+      int flatten_output_size = ori_out_size;
+      Array<Expr> fake_ops = NullValue<Array<Expr>>();
+      Array<Expr> conditions = NullValue<Array<Expr>>();
+      // if (obf_input.defined() && obf_input->IsInstance<CallNode>()) { // not handle var node input at present, obf_input should be a call
+      if (obf_input.defined() && obf_input->IsInstance<CallNode>() && selected_fops_.count(call) != 0) { // zz: add selected_fops_ condition
+        auto input_obf_attrs = obf_input.as<CallNode>()->obf_attrs;
+        Array<FloatImm> neuron_ranges = input_obf_attrs.GetAttr<Array<FloatImm>>("neuron_ranges").value_or(Array<FloatImm>());
+        Array<Integer> indices = input_obf_attrs.GetAttr<Array<Integer>>("indices").value_or(Array<Integer>());
+        for (size_t i = 0; i < neuron_ranges.size(); i++) {
+          float neuron_range = neuron_ranges[i]->value;
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode neuron_range: " << neuron_ranges[i]->value << std::endl;
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode indices: " << indices << std::endl;
+          Expr condition = GetNewCondition(obf_input, neuron_range, indices, ret_group);
+          conditions.push_back(condition);
+
+          // create fake op, currently only relu
+          Expr fake_op = GetNewFakeOp(obf_input, ori_out_size, ret_group);
+          // infer type for fake_op
+          auto fake_op_with_type = InferType(fake_op);
+          fake_op->checked_type_ = fake_op_with_type->checked_type();
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode fake_op: " << std::endl << PrettyPrint(fake_op) << std::endl;
+          fake_ops.push_back(fake_op);
+
+          int fake_op_out_size = 1;
+          for (IndexExpr dim : fake_op->checked_type().as<TensorTypeNode>()->shape) {
+            const int64_t* pval = tir::as_const_int(dim);
+            ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape " << fake_op->checked_type().as<TensorTypeNode>()->shape;
+            ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative shape" << *pval;
+            fake_op_out_size *= static_cast<size_t>(pval[0]);
+          }
+          if (fake_op_out_size > flatten_output_size) {
+            flatten_output_size = fake_op_out_size;
+          }
+        }
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode ori_out_size: " << ori_out_size << std::endl;
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode flatten_output_size: " << flatten_output_size << std::endl;
+      }
+      // =========================================================================================================
+      // [zz] step 3: construct if expr
+      // =========================================================================================================
+      // -------------------- zz: unify op output
+      // if (call->obf_attrs.defined() && call->obf_attrs.HasIntAttr("obf_op_id") && !(GetRef<Expr>(call) == last_op_)) { // zz: keep last_op_ as it is
+      if (obf_input.defined() && !(GetRef<Expr>(call) == last_op_)) { // zz: keep last_op_ as it is
+        // new_call = GetFlattenOuput(new_call, flatten_output_size, ret_group);
+
+
+        // ICHECK(fake_ops.size() <= 1); // only 1 fake op for now
+        ICHECK(conditions.size() == fake_ops.size());
+
+        if (fake_ops.size() > 0) {
+          new_call = GetFlattenOuput(new_call, flatten_output_size, ret_group);
+          for (size_t i = 0; i < fake_ops.size(); i++) {
+            fake_ops.Set(i, GetFlattenOuput(fake_ops[i], flatten_output_size, ret_group));
+          }
+
+          auto input_obf_attrs = obf_input.as<CallNode>()->obf_attrs;
+          size_t true_op_location = input_obf_attrs.GetAttr<Integer>("true_op_location").value_or(Integer(-1))->value;
+          // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode true_op_location: " << true_op_location << std::endl;
+          ICHECK(true_op_location >= 0 && true_op_location <= fake_ops.size());
+          Array<Expr> combined_ops = NullValue<Array<Expr>>();
+          size_t comb_op_idx = 0;
+          size_t fake_op_idx = 0;
+          for (size_t i = 0; i < fake_ops.size() + 1; i++) {
+            if (comb_op_idx != true_op_location) {
+              combined_ops.push_back(fake_ops[fake_op_idx]);
+              fake_op_idx++;
+            } else {
+              combined_ops.push_back(new_call);
+            }
+            comb_op_idx++;
+          }
+
+          ICHECK(combined_ops.size() == conditions.size() + 1);
+          new_call = CreateIfExpr(combined_ops, conditions , 0);
+
+          ginfo_[ret_group].has_fake_ops = true;
 
-      auto new_call = Call(call->op, new_args, call->attrs, call->type_args, call->span);
+          // zz: log op used for fake op insertion
+          // get obf_op_id from call obf_attrs
+          ICHECK(obf_input.defined() && obf_input->IsInstance<CallNode>());
+          auto obf_input_attrs = obf_input.as<CallNode>()->obf_attrs
+              .GetAttr<Integer>("obf_op_id").value_or(Integer(-1));
+          int obf_op_id = obf_input_attrs->value;
+          obf_op_ids_for_fops_.insert(obf_op_id);
+        }
+        else {
+          // zz: check whether call is in should_flat_nodes_
+          if (should_flat_nodes_.count(call) != 0) {
+            // zz: add flatten output for call
+            new_call = GetFlattenOuput(new_call, flatten_output_size, ret_group);
+          }
+        }
+      }
 
+      // =========================================================================================================
+      // [zz] step 4: group handling (unchanged)
+      // =========================================================================================================
       if (ret_group->root_ref == call) {
         // This is the root of the group
         // create the new call node.
-        return MakeNewFunction(ret_group, call->checked_type(), new_call);
+        // return MakeNewFunction(ret_group, call->checked_type(), new_call);
+        // auto ret_type = call->checked_type();
+        auto ret_type = new_call->checked_type();
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode MakeNewFunction new_call->checked_type():" << std::endl << PrettyPrint(ret_type) << std::endl;
+        // if (call->obf_attrs.defined() && call->obf_attrs.HasIntAttr("obf_op_id") && !(GetRef<Expr>(call) == last_op_)) { // zz: keep last_op_ as it is
+        //   // flat ret_type
+        //   ret_type = TensorType({Integer(ori_out_size)}, call->checked_type().as<TensorTypeNode>()->dtype);
+        // }
+        // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ CallNode MakeNewFunction input ret_type:" << std::endl << PrettyPrint(ret_type) << std::endl;
+        return MakeNewFunction(ret_group, ret_type, new_call, !ginfo_[ret_group].has_fake_ops);
+
       } else {
         // This is an intermediate node of a fused function
         // simply return the new call.
         return std::move(new_call);
       }
+
     } else {
       return ExprMutator::VisitExpr_(call);
     }
   }
 
   Expr Rewrite_(const TupleNode* tuple_node, const Expr& post) {
+    // std::cout << "[ZZDEBUG] FuseMutator Rewrite_ TupleNode (not in example):" << std::endl << PrettyPrint(GetRef<Expr>(tuple_node)) << std::endl << "[ZZDEBUG] FuseMutator Rewrite_ TupleNode end" << std::endl;
     auto* ret_group = gmap_.at(tuple_node)->FindRoot();
     if (ret_group->root_ref == tuple_node) {
       return ExprMutator::VisitExpr_(tuple_node);
@@ -439,6 +1580,7 @@ class FuseMutator : private MixedModeMutator {
   }
 
   Expr Rewrite_(const TupleGetItemNode* tuple_get, const Expr& post) {
+    // std::cout << "[ZZDEBUG FuseMutator Rewrite_ TupleGetItemNode (not in example):" << std::endl << PrettyPrint(GetRef<Expr>(tuple_get)) << std::endl << "[ZZDEBUG] FuseMutator Rewrite_ TupleGetItemNode end" << std::endl;
     auto* ret_group = gmap_.at(tuple_get)->FindRoot();
     auto new_tuple = GetNewArguments({tuple_get->tuple}, ret_group)[0];
     auto new_node = TupleGetItem(new_tuple, tuple_get->index);
@@ -456,6 +1598,7 @@ class FuseMutator : private MixedModeMutator {
   }
 
   Expr VisitExpr_(const LetNode* op) final {
+    // std::cout << "[ZZDEBUG] FuseMutator VisitExpr_ LetNode (not in example):" << std::endl << PrettyPrint(GetRef<Expr>(op)) << std::endl << "[ZZDEBUG] FuseMutator VisitExpr_ LetNode end" << std::endl;
     auto pre_visit = [this](const LetNode* op) {
       // Rely on the Memoizer to cache pre-visit values
       this->VisitExpr(op->var);
@@ -478,7 +1621,7 @@ class FuseMutator : private MixedModeMutator {
     return memo_[GetRef<Expr>(op)];
   }
 
-  Expr MakeNewFunction(GraphPartitioner::Group* group, Type ret_type, Expr body) {
+  Expr MakeNewFunction(GraphPartitioner::Group* group, Type ret_type, Expr body, bool is_primitive = true) { // zz: add is_primitive switch for kPrimitive attr
     // Quickly check special properties of the fused function.
     // A pass to check if the fused op contains only reshape ops.
     class CheckReshapeOnly : public ExprVisitor {
@@ -505,36 +1648,79 @@ class FuseMutator : private MixedModeMutator {
       bool has_call = false;
     } visitor;
 
+    // std::cout << "[ZZDEBUG] FuseMutator MakeNewFunction visitor(body) start" << std::endl;
     visitor(body);
+    // std::cout << "[ZZDEBUG] FuseMutator MakeNewFunction visitor(body) end" << std::endl;
     const GroupInfo& ginfo = ginfo_[group];
     auto func = Function(ginfo.params, body, ret_type, {});
-    func = WithAttr(std::move(func), attr::kPrimitive, tvm::Integer(visitor.has_call));
+    // std::cout << "[ZZDEBUG] FuseMutator MakeNewFunction func:" << std::endl << PrettyPrint(func) << std::endl;
+    // func = WithAttr(std::move(func), attr::kPrimitive, tvm::Integer(visitor.has_call));
+    if (is_primitive) {
+      func = WithAttr(std::move(func), attr::kPrimitive, tvm::Integer(visitor.has_call));
+    }
     // TODO(mbs): "reshape" cleanup.
     if (visitor.has_call && visitor.reshape_only) {
       func = WithAttr(std::move(func), attr::kReshapeOnly, tvm::Integer(visitor.reshape_only));
     }
-    return Call(func, ginfo.arguments, Attrs());
+    // return Call(func, ginfo.arguments, Attrs());
+    // zz: infertype needed for obf inserted ops
+    Expr new_call = Call(func, ginfo.arguments, Attrs());
+    // infer type for new_call
+    // std::cout << "[ZZDEBUG] FuseMutator MakeNewFunction InferType start" << std::endl;
+    auto new_call_with_type = InferType(new_call);
+    new_call->checked_type_ = new_call_with_type->checked_type();
+    return new_call;
   }
 
+  // zz: obf insertion can change the type of the input, ref GetNewArguments
   Array<Expr> GetNewArguments(const tvm::Array<Expr>& args,
                               GraphPartitioner::Group* current_group) {
     Array<Expr> new_args;
+    // std::cout << "[ZZDEBUG] FuseMutator GetNewArguments args size: " << args.size() << std::endl;
     for (auto arg : args) {
-      auto* arg_group = gmap_.at(arg.get())->FindRoot();
-      auto type = arg->checked_type();
-      Expr new_arg = this->Mutate(arg);
-      if (current_group != arg_group) {
-        if (!link_params_ || new_arg.as<ConstantNode>() == nullptr) {
-          Var param = ginfo_[current_group].GetOrAllocParam(new_arg, type);
-          new_args.push_back(param);
-        } else {
-          new_args.push_back(new_arg);
-        }
+      // auto* arg_group = gmap_.at(arg.get())->FindRoot();
+      // // auto type = arg->checked_type();
+      // Expr new_arg = this->Mutate(arg);
+      // // std::cout << "[ZZDEBUG] FuseMutator GetNewArguments arg:" << std::endl << PrettyPrint(arg) << std::endl;
+      // // std::cout << "[ZZDEBUG] FuseMutator GetNewArguments new_arg:" << std::endl << PrettyPrint(new_arg) << std::endl;
+      // if (!new_arg->checked_type_.defined()) // zz: idk why the type info can lost in some cases
+      //   new_arg->checked_type_ = arg->checked_type();
+      // auto type = new_arg->checked_type();  // zz: should be new_arg type
+      // if (current_group != arg_group) {
+      //   if (!link_params_ || new_arg.as<ConstantNode>() == nullptr) {
+      //     Var param = ginfo_[current_group].GetOrAllocParam(new_arg, type);
+      //     new_args.push_back(param);
+      //   } else {
+      //     new_args.push_back(new_arg);
+      //   }
+      // } else {
+      //   new_args.push_back(new_arg);
+      // }
+      new_args.push_back(GetNewArgument(arg, current_group)); // zz: move above to GetNewArgument
+    }
+    // std::cout << "[ZZDEBUG] FuseMutator GetNewArguments new_args end" << std::endl;
+    return new_args;
+  }
+
+  // zz: handle obf input
+  Expr GetNewArgument(const Expr& arg, GraphPartitioner::Group* current_group) {
+    auto* arg_group = gmap_.at(arg.get())->FindRoot();
+    Expr new_arg = this->Mutate(arg);
+    if (!new_arg->checked_type_.defined()) // zz: idk why the type info can lost in some cases, not needed now
+      new_arg->checked_type_ = arg->checked_type();
+    auto type = new_arg->checked_type();
+    if (current_group != arg_group) {
+      if (!link_params_ || new_arg.as<ConstantNode>() == nullptr) {
+        Var param = ginfo_[current_group].GetOrAllocParam(new_arg, type);
+        // param->checked_type_ = type; // zz: set type for param, not needed now
+        return param;
       } else {
-        new_args.push_back(new_arg);
+        return new_arg;
       }
+    } else {
+      // std::cout << "[ZZDEBUG] FuseMutator GetNewArgument new_arg" << std::endl << PrettyPrint(new_arg) << std::endl << "[ZZDEBUG] FuseMutator GetNewArgument new_arg end" << std::endl;
+      return new_arg;
     }
-    return new_args;
   }
 
   // Debug function, dump the group assignment in text.
@@ -549,17 +1735,41 @@ class FuseMutator : private MixedModeMutator {
     });
     LOG(INFO) << "Dump of group info:\n" << text;
   }
+
+  // zz: for debuging purpose
+  std::string DebugDumpGroupStr(const Expr& body) {
+    std::string text = AsText(body, false, [this](const ObjectRef& expr) -> std::string {
+      auto it = gmap_.find(expr.get());
+      if (it == gmap_.end()) return "";
+      std::ostringstream os;
+      auto* group = it->second->FindRoot();
+      os << " /* group=" << group << " */";
+      return os.str();
+    });
+    return text;
+  }
 };
 
 Expr FuseOps(const Expr& expr, int fuse_opt_level, size_t max_fuse_depth, size_t max_function_args,
-             bool link_params, const IRModule& module) {
-  return FuseMutator(fuse_opt_level, max_fuse_depth, max_function_args, link_params)
+             bool link_params, const IRModule& module, bool enable_flex_fuse, bool enable_fakeop_insert) { // zz: add obf switch
+  // std::cout << "[ZZDEBUG] FuseOps: fuse_opt_level = " << fuse_opt_level << ", max_fuse_depth = " << max_fuse_depth << ", max_function_args = " << max_function_args << std::endl;
+  return FuseMutator(fuse_opt_level, max_fuse_depth, max_function_args, link_params, enable_flex_fuse, enable_fakeop_insert)
       .Transform(expr);
 }
 
+// zz: for random fuse, ref FuseOps(const Expr& expr, ...)
+Expr FuseOpsObf(const Expr& expr, int fuse_opt_level, size_t max_fuse_depth, size_t max_function_args,
+            //  bool link_params, const IRModule& module) {
+            //  bool link_params, const IRModule& module, bool skip_fuse) {
+             bool link_params, const IRModule& module, bool enable_flex_fuse, bool enable_fakeop_insert) { // zz: add obf switch
+  // return FuseMutator(fuse_opt_level, max_fuse_depth, max_function_args, link_params, skip_fuse)
+  return FuseMutator(fuse_opt_level, max_fuse_depth, max_function_args, link_params, enable_flex_fuse, enable_fakeop_insert)
+      .TransformObf(expr);
+}
+
 namespace transform {
 
-Pass FuseOps(int fuse_opt_level) {
+Pass FuseOps(int fuse_opt_level, bool enable_flex_fuse, bool enable_fakeop_insert) { // zz: add obf switch
   runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
       [=](Function f, IRModule m, PassContext pc) {
         bool link_params = false;
@@ -577,13 +1787,359 @@ Pass FuseOps(int fuse_opt_level) {
                 ? target->GetAttr<Integer>("max_function_args", Integer(0)).value().IntValue()
                 : 0;
         return Downcast<Function>(FuseOps(f, opt_level, max_fuse_depth.value().IntValue(),
-                                          max_function_args, link_params, m));
+                                          max_function_args, link_params, m, enable_flex_fuse, enable_fakeop_insert)); // zz: add obf switch
       };
   return CreateFunctionPass(pass_func, 0, "FuseOps", {"InferType"});
 }
 
 TVM_REGISTER_GLOBAL("relay._transform.FuseOps").set_body_typed(FuseOps);
 
+// zz: for random fuse, ref FuseOps(int fuse_opt_level)
+Pass FuseOpsObf(int fuse_opt_level, bool enable_flex_fuse, bool enable_fakeop_insert) { // zz: add obf switch
+  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
+      [=](Function f, IRModule m, PassContext pc) {
+        // std::cout << "[ZZDEBUG] FuseOpsObf at func :" << std::endl << PrettyPrint(f) << std::endl;
+        bool link_params = false;
+        Executor executor =
+            m->GetAttr<Executor>(tvm::attr::kExecutor).value_or(NullValue<Executor>());
+        link_params = executor.defined()
+                          ? executor->attrs.GetAttr<Bool>("link-params").value_or(Bool(link_params))
+                          : link_params;
+        link_params = pc->GetConfig("relay.FuseOps.link_params", Bool(link_params)).value();
+        int opt_level = fuse_opt_level == -1 ? pc->opt_level : fuse_opt_level;
+        auto max_fuse_depth = pc->GetConfig("relay.FuseOps.max_depth", Integer(kMaxFusedOps));
+        auto target = Target::Current();
+        size_t max_function_args =
+            (target.defined())
+                ? target->GetAttr<Integer>("max_function_args", Integer(0)).value().IntValue()
+                : 0;
+        return Downcast<Function>(FuseOpsObf(f, opt_level, max_fuse_depth.value().IntValue(),
+                                          // max_function_args, link_params, m, skip_fuse));
+                                          max_function_args, link_params, m, enable_flex_fuse, enable_fakeop_insert));// zz: add obf switch
+      };
+  return CreateFunctionPass(pass_func, 0, "FuseOpsObf", {"InferType"});
+}
+
+TVM_REGISTER_GLOBAL("relay._transform.FuseOpsObf").set_body_typed(FuseOpsObf);
+
+Expr AnnotateObfAttrs(const Expr& expr) {
+  class AnnotateObfAttrsMutator : public ExprMutator {
+   public:
+    AnnotateObfAttrsMutator() {
+      transform::PassContext pass_ctx = transform::PassContext::Current();
+      selected_neuron_indices = pass_ctx->GetConfig("dnnobf.selected_neuron_indices", Array<Array<Integer>>({})).value();
+      selected_neuron_ranges = pass_ctx->GetConfig("dnnobf.selected_neuron_ranges", Array<Array<FloatImm>>({})).value();
+      // std::cout << "[ZZDEBUG] AnnotateObfAttrsMutator selected_neuron_indices:" << std::endl << PrettyPrint(selected_neuron_indices) << std::endl;
+      // std::cout << "[ZZDEBUG] AnnotateObfAttrsMutator selected_neuron_ranges:" << std::endl << PrettyPrint(selected_neuron_ranges) << std::endl;
+    }
+    Expr VisitExpr_(const CallNode* call_node) final {
+      auto new_op = this->Mutate(call_node->op);
+
+      tvm::Array<Type> ty_args;
+      ty_args.reserve(call_node->type_args.size());
+
+      for (auto ty_arg : call_node->type_args) {
+        auto new_ty_arg = this->VisitType(ty_arg);
+        ty_args.push_back(new_ty_arg);
+      }
+
+      tvm::Array<Expr> call_args;
+      call_args.reserve(call_node->args.size());
+      for (auto arg : call_node->args) {
+        auto new_arg = this->Mutate(arg);
+        call_args.push_back(new_arg);
+      }
+
+      Expr before_anotate = WithFields(GetRef<Call>(call_node), new_op, call_args, {}, ty_args);
+      // only annotate if call_node output float32 tensor
+      if (!call_node->checked_type().as<TensorTypeNode>()) {
+        return before_anotate;
+      }
+      if (call_node->checked_type().as<TensorTypeNode>()->dtype != DataType::Float(32)) {
+        return before_anotate;
+      }
+      // // only annotate if call_node inpputs are float32 tensor
+      // for (auto arg : call_node->args) {
+      //   if (arg->checked_type().as<TensorTypeNode>()) {
+      //     if (arg->checked_type().as<TensorTypeNode>()->dtype != DataType::Float(32)) {
+      //       return before_anotate;
+      //     }
+      //   } else {
+      //     // if arg is not a tensor type, skip annotation
+      //     return before_anotate;
+      //   }
+      // }
+
+      Map<String, ObjectRef> dict;
+      // get call_node shape
+      auto ori_out_shape = call_node->checked_type().as<TensorTypeNode>()->shape;
+      int dim = ori_out_shape.size();
+
+      if (selected_neuron_indices.size() == 0) {
+        Array<Integer> indices;
+        for (int i = 0; i < dim; i++) {
+          indices.push_back(Integer(0));
+        }
+
+        Array<FloatImm> neuron_ranges;
+        neuron_ranges.push_back(FloatImm(DataType::Float(32), -1000.5 - interesting_exprs_id));
+        neuron_ranges.push_back(FloatImm(DataType::Float(32), 1000.5 + interesting_exprs_id));
+
+        dict.Set("obf_op_id", Integer(interesting_exprs_id++));
+        dict.Set("indices", indices);
+        dict.Set("neuron_ranges", neuron_ranges);
+        dict.Set("true_op_location", Integer(1));
+        DictAttrs obf_attrs(dict);
+
+        return WithFields(GetRef<Call>(call_node), new_op, call_args, {}, ty_args, {}, {}, obf_attrs);
+      } else {
+        Array<Integer> indices = selected_neuron_indices[interesting_exprs_id];
+        float low = selected_neuron_ranges[interesting_exprs_id][0]->value;
+        float high = selected_neuron_ranges[interesting_exprs_id][1]->value;
+        int true_op_location = rand() % 3; // TODO: make insert_width configurable
+        Array<FloatImm> neuron_ranges;
+        if (true_op_location == 0) {
+          neuron_ranges.push_back(FloatImm(DataType::Float(32), high));
+          float range_length = high - low;
+          float increse_length = (static_cast <float> (rand()) / static_cast <float> (RAND_MAX)) * range_length + 0.5 * range_length;
+          neuron_ranges.push_back(FloatImm(DataType::Float(32), high + increse_length));
+        } else if (true_op_location == 1) {
+          neuron_ranges.push_back(FloatImm(DataType::Float(32), low));
+          neuron_ranges.push_back(FloatImm(DataType::Float(32), high));
+        } else { // true_op_location == 2
+          float range_length = high - low;
+          float decrese_length = (static_cast <float> (rand()) / static_cast <float> (RAND_MAX)) * range_length + 0.5 * range_length;
+          neuron_ranges.push_back(FloatImm(DataType::Float(32), low - decrese_length));
+          neuron_ranges.push_back(FloatImm(DataType::Float(32), low));
+        }
+
+        dict.Set("obf_op_id", Integer(interesting_exprs_id++));
+        dict.Set("indices", indices);
+        dict.Set("neuron_ranges", neuron_ranges);
+        dict.Set("true_op_location", Integer(true_op_location));
+        DictAttrs obf_attrs(dict);
+
+        return WithFields(GetRef<Call>(call_node), new_op, call_args, {}, ty_args, {}, {}, obf_attrs);
+      }
+    }
+
+    int interesting_exprs_id = 0;
+    Array<Array<Integer>> selected_neuron_indices = {};
+    Array<Array<FloatImm>> selected_neuron_ranges = {};
+  };
+
+  // return AnnotateObfAttrsMutator().Mutate(expr);
+  AnnotateObfAttrsMutator aomutator;
+  Expr result = aomutator.Mutate(expr);
+  // std::cout << "[ZZDEBUG] AnnotateObfAttrs # of interesting ops = " << aomutator.interesting_exprs_id << std::endl;
+  ICHECK(aomutator.selected_neuron_indices.size() == aomutator.selected_neuron_ranges.size());
+  if (aomutator.selected_neuron_indices.size() > 0) {
+    ICHECK(aomutator.interesting_exprs_id == aomutator.selected_neuron_indices.size());
+  }
+  return result;
+}
+
+Pass AnnotateObfAttrs() {
+  // function pass
+  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
+      [=](Function f, IRModule m, PassContext pc) {
+    return Downcast<Function>(AnnotateObfAttrs(f)); // zz: add obf switch
+  };
+
+  return CreateFunctionPass(pass_func, 0, "AnnotateObfAttrs", {"InferType"});
+}
+
+/*
+@relay.transform.function_pass(opt_level=1)
+def ftransform(func: relay.Function, mod, ctx):
+    global interesting_exprs_bounds
+    global interesting_exprs
+    # global out_idx_2_op_name
+    global op_names
+    # global op_index
+
+    class InterestingExprCollector(relay.ExprVisitor):
+
+        def __init__(self):
+            super().__init__()
+            self.skipped_exprs = set()
+            self.op_index = 0
+
+        def visit_call(self, call):
+            # Goal: Instrument conv2d and dense layers
+            op_name = get_op_name(call.op)
+            # print(f'******************** python Visiting {op_name}', flush=True)
+            # print(f'******************** python visit_call: get_type={get_type(call)}, get_dtype={get_dtype(call)}, get_shape={get_shape(call)}', flush=True)
+            super().visit_call(call) # recursive relay.ExprVisitor
+            # return call
+            # if not isinstance(call, r.Tuple):
+            #     interesting_exprs.append(call)
+            if op_name != "nn.batch_norm":
+                interesting_exprs.append(call)
+                op_names.append(op_name+ f"_{self.op_index}")
+                # out_idx_2_op_name[self.op_index] = op_name
+                self.op_index += 1
+    InterestingExprCollector().visit(func.body)
+    for i, e in enumerate(interesting_exprs):
+        print(f'******************** interesting_exprs[{i}]: {e}', flush=True)
+        shape = get_shape(e)
+        interesting_exprs_bounds.append({})
+        interesting_exprs_bounds[i]["low"] = np.zeros(shape)
+        interesting_exprs_bounds[i]["high"] = np.zeros(shape)
+        # print("interesting_exprs[{}] bounds shape: {}".format(i, shape), flush=True)
+
+    # print("interesting_exprs: {}".format([get_op_name(e.op) for e in interesting_exprs]), flush=True)
+    for i, e in enumerate(interesting_exprs):
+        print(f'******************** interesting_exprs[{i}]: get_type={get_type(e)}, get_dtype={get_dtype(e)}, get_shape={get_shape(e)}', flush=True)
+
+    layer_ret_types = [func.ret_type]
+    class Mutator(r.ExprMutator):
+        def visit(self, expr: relay.Expr):
+            for e in interesting_exprs:
+                shape = get_shape(e)
+                dtype = get_dtype(e)
+                layer_ret_types.append(r.TensorType(shape, dtype=dtype))
+
+            out_tuple = r.Tuple([expr] + interesting_exprs)
+            return out_tuple
+    new_body = Mutator().visit(func.body)
+    rettype = r.TupleType(layer_ret_types)
+
+    return relay.Function(func.params, new_body, rettype, func.type_params, func.attrs)
+
+*/
+
+Expr InterestingExprExaminer(const Expr& expr) {
+  class InterestingExprCollector: public ExprVisitor {
+   public:
+    InterestingExprCollector() {
+      interesting_exprs_id = 0;
+    }
+
+    void VisitExpr_(const CallNode* call_node) final {
+      // do ExprVisitor visit
+      ExprVisitor::VisitExpr_(call_node);
+      // this->VisitSpan(call_node->span);
+      // this->VisitExpr(call_node->op);
+
+      // for (auto ty_arg : call_node->type_args) {
+      //   this->VisitType(ty_arg);
+      // }
+
+      // for (auto arg : call_node->args) {
+      //   this->VisitExpr(arg);
+      // }
+
+      // only collect call_node if it is a float32 tensor output
+      if (!call_node->checked_type().as<TensorTypeNode>()) {
+        return;
+      }
+      if (call_node->checked_type().as<TensorTypeNode>()->dtype != DataType::Float(32)) {
+        return;
+      }
+
+      interesting_exprs.push_back(GetRef<Call>(call_node));
+      interesting_exprs_types.push_back(call_node->checked_type());
+      interesting_exprs_id++;
+    }
+
+    int interesting_exprs_id;
+    Array<Expr> interesting_exprs; // collect interesting exprs
+    // array of expr type
+    Array<Type> interesting_exprs_types; // collect interesting exprs types
+  };
+
+  // check expr is a function
+  if (!expr->IsInstance<FunctionNode>()) {
+    LOG(FATAL) << "InterestingExprExaminer expects a Function, but got: " << expr->GetTypeKey();
+  }
+  // convert to Function
+  Function func = Downcast<Function>(expr);
+
+
+  InterestingExprCollector collector;
+  collector(expr);
+
+  // make relay::Tupple for interesting exprs
+  Expr interesting_exprs_tuple = Tuple(collector.interesting_exprs);
+  Type interesting_exprs_tuple_type = TupleType(collector.interesting_exprs_types);
+  // std::cout << "[ZZDEBUG] InterestingExprExaminer interesting_exprs size = " << collector.interesting_exprs.size() << std::endl;
+  // std::cout << "[ZZDEBUG] InterestingExprExaminer interesting_exprs_tuple: " << PrettyPrint(interesting_exprs_tuple) << std::endl;
+
+  return Function(func->params, interesting_exprs_tuple, interesting_exprs_tuple_type, func->type_params);
+
+
+
+  // class InterestingExprCollectorMutator : public ExprMutator {
+  //  public:
+  //   Expr VisitExpr_(const CallNode* call_node) final {
+  //     auto new_op = this->Mutate(call_node->op);
+
+  //     tvm::Array<Type> ty_args;
+  //     ty_args.reserve(call_node->type_args.size());
+
+  //     for (auto ty_arg : call_node->type_args) {
+  //       auto new_ty_arg = this->VisitType(ty_arg);
+  //       ty_args.push_back(new_ty_arg);
+  //     }
+
+  //     tvm::Array<Expr> call_args;
+  //     call_args.reserve(call_node->args.size());
+  //     for (auto arg : call_node->args) {
+  //       auto new_arg = this->Mutate(arg);
+  //       call_args.push_back(new_arg);
+  //     }
+
+  //     Expr before_anotate = WithFields(GetRef<Call>(call_node), new_op, call_args, {}, ty_args);
+  //     // only annotate if call_node output float32 tensor
+  //     if (!call_node->checked_type().as<TensorTypeNode>()) {
+  //       return before_anotate;
+  //     }
+  //     if (call_node->checked_type().as<TensorTypeNode>()->dtype != DataType::Float(32)) {
+  //       return before_anotate;
+  //     }
+
+  //     Map<String, ObjectRef> dict;
+  //     // get call_node shape
+  //     auto ori_out_shape = call_node->checked_type().as<TensorTypeNode>()->shape;
+  //     int dim = ori_out_shape.size();
+  //     Array<Integer> indices;
+  //     for (int i = 0; i < dim; i++) {
+  //       indices.push_back(Integer(0));
+  //     }
+
+  //     Array<FloatImm> neuron_ranges;
+  //     neuron_ranges.push_back(FloatImm(DataType::Float(32), -1000.5 - interesting_exprs_id));
+  //     neuron_ranges.push_back(FloatImm(DataType::Float(32), 1000.5 + interesting_exprs_id));
+
+  //     dict.Set("obf_op_id", Integer(interesting_exprs_id++));
+  //     dict.Set("indices", indices);
+  //     dict.Set("neuron_ranges", neuron_ranges);
+  //     dict.Set("true_op_location", Integer(1));
+  //     DictAttrs obf_attrs(dict);
+
+  //     return WithFields(GetRef<Call>(call_node), new_op, call_args, {}, ty_args, {}, {}, obf_attrs);
+  //   }
+
+  //   int interesting_exprs_id = 0;
+  // };
+
+  // return InterestingExprCollectorMutator().Mutate(expr);
+
+}
+
+Pass InterestingExprExaminer() {
+  // function pass
+  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
+      [=](Function f, IRModule m, PassContext pc) {
+    return Downcast<Function>(InterestingExprExaminer(f)); // zz: add obf switch
+  };
+
+  return CreateFunctionPass(pass_func, 0, "InterestingExprExaminer", {"InferType"});
+}
+
+// TVM_REGISTER_GLOBAL("relay._transform.FuseOpsObf").set_body_typed(FuseOpsObf);
+
 }  // namespace transform
 
 }  // namespace relay
diff --git a/src/relay/transforms/remove_standalone_reshapes.cc b/src/relay/transforms/remove_standalone_reshapes.cc
index 063060b3e..09ab780d0 100644
--- a/src/relay/transforms/remove_standalone_reshapes.cc
+++ b/src/relay/transforms/remove_standalone_reshapes.cc
@@ -110,6 +110,29 @@ Pass RemoveStandaloneReshapes() {
   return tvm::transform::CreateModulePass(pass_func, 0, "RemoveStandaloneReshapes", {});
 }
 
+// zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+Pass RemoveStandaloneReshapesWithMainFuncName(String main_func_name) {
+  auto pass_func = [=](IRModule mod, const PassContext& pass_ctx) {
+    VLOG(1) << "RemoveStandaloneReshapes before:" << std::endl << PrettyPrint(mod);
+    RemoveStandaloneReshapesMutator remove_reshapes_mutator(mod);
+    Function main_func = Downcast<Function>(mod->Lookup(main_func_name));
+    Expr new_main_body = remove_reshapes_mutator.VisitExpr(main_func->body);
+    if (!new_main_body.same_as(main_func->body)) {
+      auto main_var = mod->GetGlobalVar(main_func_name);
+      auto new_main_func = Function(main_func->params, new_main_body, main_func->ret_type,
+                                    main_func->type_params, main_func->attrs);
+      mod->Update(main_var, new_main_func);
+    }
+    Array<runtime::String> entry_functions{main_func_name};
+    mod = RemoveUnusedFunctions(entry_functions)(mod);
+
+    VLOG(1) << "RemoveStandaloneReshapes after:" << std::endl << PrettyPrint(mod);
+    return mod;
+  };
+  return tvm::transform::CreateModulePass(pass_func, 0, "RemoveStandaloneReshapes", {});
+}
+
+
 TVM_REGISTER_GLOBAL("relay._transform.RemoveStandaloneReshapes")
     .set_body_typed(RemoveStandaloneReshapes);
 
diff --git a/src/relay/transforms/to_a_normal_form.cc b/src/relay/transforms/to_a_normal_form.cc
index 8319726b7..c88e6cd19 100644
--- a/src/relay/transforms/to_a_normal_form.cc
+++ b/src/relay/transforms/to_a_normal_form.cc
@@ -256,7 +256,8 @@ class Fill : ExprFunctor<Expr(const Expr&, const Var&)>, private transform::Lexi
     for (const auto& a : c->args) {
       args.push_back(VisitExpr(a));
     }
-    return Compound(e, Call(VisitExpr(c->op), args, c->attrs, c->type_args), v);
+    // return Compound(e, Call(VisitExpr(c->op), args, c->attrs, c->type_args), v);
+    return Compound(e, Call(VisitExpr(c->op), args, c->attrs, c->type_args, Span(), c->obf_attrs), v); // zz: add obf_attrs
   }
 
   Expr VisitExpr_(const TupleNode* tuple_node, const Var& v) final {
diff --git a/src/target/metadata.h b/src/target/metadata.h
index b761f7ff2..6ba153478 100644
--- a/src/target/metadata.h
+++ b/src/target/metadata.h
@@ -153,6 +153,7 @@ class InMemoryMetadataNode : public ::tvm::target::metadata::VisitableMetadataNo
     storage_.constant_pools = constant_pools_.get();
     storage_.num_constant_pools = constant_pools.size();
     for (size_t i = 0; i < constant_pools.size(); ++i) {
+      std::cout << "[ZZDEBUG] InMemoryMetadataNode constant_pools_ " << i << ": " << constant_pools[i]->name_hint << std::endl;
       constant_pools_.get()[i].name_hint = constant_pools[i]->name_hint.c_str();
       constant_pools_.get()[i].byte_offset = constant_pools[i]->byte_offset.IntValue();
 
@@ -160,11 +161,30 @@ class InMemoryMetadataNode : public ::tvm::target::metadata::VisitableMetadataNo
       dmlc::MemoryStringStream stream(&bytes);
       auto data = constant_pools[i]->data;
       data.Save(&stream);
+      // zz: Save() called above is added with debug info
+      std::cout << "  (zz) allocated bytes.size(): " << bytes.size() << std::endl;
+
       // Allocated mem freed in destructor
       constant_pools_.get()[i].data_len = bytes.size();
       char* a = reinterpret_cast<char*>(malloc(bytes.size()));
       constant_pools_.get()[i].data_bytes = a;
       memcpy(a, bytes.c_str(), bytes.size());
+
+
+      // ----------------------------------------------
+      // zz: shape, tsize, tensor byte_size debug info
+      // ----------------------------------------------
+      int tsize = 1;
+      for (int j = 0; j < data->ndim; ++j) {
+        std::cout << "  (zz) shape[" << j << "]: " << data->shape[j] << std::endl;
+        tsize *= data->shape[j];
+      }
+      std::cout << "  (zz) tsize: " << tsize << std::endl;
+      int byte_size = tsize * ((data->dtype.bits * data->dtype.lanes + 7) / 8);
+      std::cout << "  (zz) tensor byte_size: " << byte_size << std::endl;
+      // ----------------------------------------------------
+      // zz: end - shape, tsize, tensor byte_size debug info
+      // ----------------------------------------------------
     }
   }
 
diff --git a/src/target/metadata_module.cc b/src/target/metadata_module.cc
index c8c099171..443c303b6 100644
--- a/src/target/metadata_module.cc
+++ b/src/target/metadata_module.cc
@@ -46,6 +46,7 @@ static runtime::Module CreateCrtMetadataModule(
     Array<runtime::Module> crt_exportable_modules,
     const std::unordered_map<std::string, runtime::NDArray>& const_var_ndarray) {
   if (!non_crt_exportable_modules.empty()) {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateCrtMetadataModule !non_crt_exportable_modules.empty()" << std::endl;
     std::string non_exportable_modules;
     for (unsigned int i = 0; i < non_crt_exportable_modules.size(); i++) {
       if (i > 0) {
@@ -65,6 +66,7 @@ static runtime::Module CreateCrtMetadataModule(
   }
 
   if (target->kind->name == "c") {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateCrtMetadataModule target->kind->name == c" << std::endl;
     runtime::metadata::Metadata aot_metadata;
     if (executor->GetAttr<String>("interface-api", tvm::String("packed")) == "packed") {
       aot_metadata = ConvertMetaData(metadata);
@@ -74,6 +76,9 @@ static runtime::Module CreateCrtMetadataModule(
     target_module = CreateCSourceCrtMetadataModule(crt_exportable_modules, target, runtime,
                                                    metadata, aot_metadata);
   } else if (target->kind->name == "llvm") {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateCrtMetadataModule target->kind->name == llvm" << std::endl;
+    runtime::metadata::Metadata aot_metadata; // zz: to call ConvertMetaData for now. TODO: remove this
+    aot_metadata = ConvertMetaData(metadata); // zz: added to call SaveData for params saving for now. TODO: remove this
 #ifdef TVM_LLVM_VERSION
     crt_exportable_modules.push_back(target_module);
     target_module = CreateLLVMCrtMetadataModule(crt_exportable_modules, target, runtime);
@@ -88,6 +93,7 @@ static runtime::Module CreateCrtMetadataModule(
 // TODO(areusch,masahi): Unify metadata representation and remove the need for this function
 static runtime::metadata::Metadata ConvertMetaData(
     relay::backend::ExecutorCodegenMetadata metadata) {
+  std::cout << "[ZZDEBUG] tvm::codegen::ConvertMetaData" << std::endl;
   ICHECK(metadata.defined());
   ICHECK_NOTNULL(metadata->pool_inputs);
 
@@ -124,6 +130,8 @@ static runtime::metadata::Metadata ConvertMetaData(
   }
 
   std::vector<ConstantInfo> consts;
+  int num_constants = 0; // zz: should have only 1 constant pool
+  std::vector<std::pair<size_t, size_t>> constant_idx_byte_offsets; // zz: sort later
   for (const auto& kv : metadata->pool_inputs.value()) {
     const auto& api = kv.second;
     if (const auto* pi = api->pool_info.as<ConstantPoolInfoNode>()) {
@@ -131,12 +139,103 @@ static runtime::metadata::Metadata ConvertMetaData(
         for (const auto ci : pi->constant_info_array) {
           consts.emplace_back(ci->name_hint, ci->byte_offset, ci->data);
         }
+
+        //----------------------------------------------------
+        // zz: dump constant info (stdout & binary file)
+        //----------------------------------------------------
+        num_constants++;
+        std::cout << "[ZZDEBUG] tvm::codegen::ConvertMetaData constant_info_array:" << std::endl;
+        for (size_t i = 0; i < pi->constant_info_array.size(); i++) {
+          auto ci = pi->constant_info_array[i];
+          std::cout << "  (zz) constant_info_array " << i << ": " << ci->name_hint << ", byte_offset=" << ci->byte_offset << std::endl;
+          constant_idx_byte_offsets.emplace_back(std::make_pair(i, ci->byte_offset.IntValue()));
+        }
+
+        // zz: group constant_info_array to a single buffer for aot direct call
+        // sort constant_idx_byte_offsets using byte_offset increasing order
+        std::sort(constant_idx_byte_offsets.begin(), constant_idx_byte_offsets.end(),
+                  [](const std::pair<size_t, size_t>& a, const std::pair<size_t, size_t>& b) {
+                    return a.second < b.second;
+                  });
+
+        std::string bytes;
+        dmlc::MemoryStringStream stream(&bytes);
+        size_t total_data_bytes = 0;
+        size_t total_const_buf_size = 0;
+        for(size_t i = 0; i < constant_idx_byte_offsets.size(); i++) {
+          size_t idx = constant_idx_byte_offsets[i].first;
+          auto ci = pi->constant_info_array[idx];
+          auto data = ci->data;
+          size_t data_offset = ci->byte_offset.IntValue();
+          if (total_const_buf_size < data_offset) {
+            // zz: fill gap with 0
+            size_t gap_size = data_offset - total_const_buf_size;
+            std::vector<uint8_t> gap(gap_size, 0);
+            stream.Write(gap.data(), gap.size());
+            total_const_buf_size += gap_size;
+          }
+          // data.SaveData(&stream);
+          size_t data_size = data.SaveData(&stream);
+          total_data_bytes += data_size;
+          total_const_buf_size += data_size;
+          std::cout << "  (zz) sorted constant_info_array " << i << ": " << ci->name_hint << ", byte_offset (include gap) = " << ci->byte_offset << ", current total_data_bytes (data size) = " << total_data_bytes << std::endl;
+        }
+        // wirte stream to a binary file
+        std::ofstream ofs("const.bin", std::ios::binary);
+        ofs.write(bytes.c_str(), bytes.size());
+        //----------------------------------------------------
+        // zz: dump constant info (stdout & binary file) end
+        //----------------------------------------------------
       }
     }
   }
+  ICHECK(num_constants <= 1) << "[ZZDEBUG] At most one constant pool (global) is supported";
   auto n = make_object<target::metadata::InMemoryMetadataNode>(
       runtime::metadata::kMetadataVersion, inputs, outputs, pools, consts, metadata->mod_name);
 
+  // =================================================================================
+  // zz: dump metadata info
+  // =================================================================================
+  std::cout << "[ZZDEBUG] tvm::codegen::ConvertMetaData dump metadata:" << std::endl;
+  const struct ::TVMMetadata* md = n->get_metadata();
+  printf("\tmod_name=%s\n", md->mod_name);
+  printf("\tversion=%" PRId64 "\n", md->version);
+  printf("\tnum_inputs=%" PRId64 "\n", md->num_inputs);
+  printf("\tnum_outputs=%" PRId64 "\n", md->num_outputs);
+  printf("\tnum_workspace_pools=%" PRId64 "\n", md->num_workspace_pools);
+  printf("\tnum_constant_pools=%" PRId64 "\n", md->num_constant_pools);
+
+
+  for (int i = 0; i < md->num_inputs; ++i) {
+    printf("\tinput[%d]: %s\n", i, md->inputs[i].name);
+    for (int si = 0; si < md->inputs[i].num_shape; si++) {
+      printf("\t\t%ld,\n", md->inputs[i].shape[si]);
+    }
+  }
+
+  for (int i = 0; i < md->num_outputs; ++i) {
+    printf("\toutput[%d]: %s\n", i, md->outputs[i].name);
+    for (int si = 0; si < md->outputs[i].num_shape; si++) {
+      printf("\t\t%ld,\n", md->outputs[i].shape[si]);
+    }
+  }
+
+  for (int i = 0; i < md->num_workspace_pools; ++i) {
+    printf("\tworkspace_pools[%d]: %s\n", i, md->workspace_pools[i].name);
+    for (int si = 0; si < md->workspace_pools[i].num_shape; si++) {
+      printf("\t\t%ld,\n", md->workspace_pools[i].shape[si]);
+    }
+  }
+
+  for (int i = 0; i < md->num_constant_pools; ++i) {
+    printf("\tconstant_pools[%d]: %s\n", i, md->constant_pools[i].name_hint);
+    printf("\t\tbyte_offset=%ld, data_len=%ld, data_bytes=%p\n", md->constant_pools[i].byte_offset, md->constant_pools[i].data_len, md->constant_pools[i].data_bytes);
+  }
+  std::cout << "[ZZDEBUG] tvm::codegen::ConvertMetaData dump metadata end." << std::endl;
+  // =================================================================================
+  // zz: (end) dump metadata info
+  // =================================================================================
+
   return runtime::metadata::Metadata(std::move(n));
 }
 
@@ -148,6 +247,7 @@ static runtime::Module CreateCppMetadataModule(
     Array<runtime::Module> crt_exportable_modules,
     const std::unordered_map<std::string, runtime::NDArray>& const_var_ndarray) {
   if (!non_crt_exportable_modules.empty()) {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateCppMetadataModule !non_crt_exportable_modules.empty()" << std::endl;
     runtime::Module const_loader_mod =
         runtime::ConstLoaderModuleCreate(const_var_ndarray, const_vars_by_symbol);
     const_loader_mod.Import(target_module);
@@ -158,15 +258,19 @@ static runtime::Module CreateCppMetadataModule(
   }
 
   if (metadata.defined()) {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateCppMetadataModule metadata.defined()" << std::endl;
     runtime::metadata::Metadata runtime_metadata = ConvertMetaData(metadata);
 
     if (metadata->executor == runtime::kTvmExecutorAot && runtime->name == relay::kTvmRuntimeCpp) {
+      std::cout << "[ZZDEBUG] tvm::codegen::CreateCppMetadataModule metadata->executor == runtime::kTvmExecutorAot && runtime->name == relay::kTvmRuntimeCpp" << std::endl;
       if (target->kind->name == "c") {
+        std::cout << "[ZZDEBUG] tvm::codegen::CreateCppMetadataModule target->kind->name == c" << std::endl;
         auto metadata_module = CreateCSourceCppMetadataModule(runtime_metadata);
         metadata_module->Import(target_module);
         target_module = metadata_module;
 #ifdef TVM_LLVM_VERSION  // defining TVM_LLVM_VERSION indicates TVM was compiled with USE_LLVM ON.
       } else if (target->kind->name == "llvm") {
+        std::cout << "[ZZDEBUG] tvm::codegen::CreateCppMetadataModule target->kind->name == llvm" << std::endl;
         auto metadata_module = CreateLLVMCppMetadataModule(runtime_metadata, target, runtime);
         metadata_module->Import(target_module);
         target_module = metadata_module;
@@ -175,6 +279,7 @@ static runtime::Module CreateCppMetadataModule(
         CHECK(false) << "Don't know how to create MetadataModule for target type " << target->str();
       }
     }
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateCppMetadataModule metadata handling end." << std::endl;
   }
 
   return target_module;
@@ -196,6 +301,34 @@ runtime::Module CreateMetadataModule(
     tvm::runtime::Module target_module, const Array<runtime::Module>& ext_modules, Target target,
     tvm::relay::Runtime runtime, tvm::relay::Executor executor,
     relay::backend::ExecutorCodegenMetadata metadata) {
+
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateMetadataModule const_var_ndarray:" << std::endl;
+    int i = 0;
+    int total_const_byte_size = 0;
+    for (const auto& kv : const_var_ndarray) {
+  //  * static inline size_t GetDataSize(const DLTensor* t) {
+  //  *   size_t size = 1;
+  //  *   for (tvm_index_t i = 0; i < t->ndim; ++i) {
+  //  *     size *= t->shape[i];
+  //  *   }
+  //  *   size *= (t->dtype.bits * t->dtype.lanes + 7) / 8;
+  //  *   return size;
+  //  * }
+      std::cout << "    const_var_ndarray " << i << ": " << kv.first << std::endl;
+      const DLTensor* t = kv.second.operator->();
+      int tsize = 1;
+      for (int j = 0; j < t->ndim; j++) {
+        std::cout << "        shape[" << j << "]: " << t->shape[j] << std::endl;
+        tsize *= t->shape[j];
+      }
+      std::cout << "        tsize: " << tsize << std::endl;
+      int byte_size = tsize * ((t->dtype.bits * t->dtype.lanes + 7) / 8);
+      std::cout << "        bytes: " << byte_size << std::endl;
+      total_const_byte_size += byte_size;
+      i++;
+    }
+    std::cout << "total_const_byte_size: " << total_const_byte_size << std::endl;
+
   // Here we split modules into two groups:
   //  1. Those modules which can be exported to C-runtime. These are DSO-exportable
   //     (i.e. llvm or c) modules which return nothing from get_const_vars().
@@ -208,6 +341,7 @@ runtime::Module CreateMetadataModule(
   // Wrap all submodules in the initialization wrapper.
   std::unordered_map<std::string, std::vector<std::string>> const_vars_by_symbol;
   for (tvm::runtime::Module mod : ext_modules) {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateMetadataModule ext_modules" << std::endl;
     auto pf_sym = mod.GetFunction("get_symbol");
     auto pf_var = mod.GetFunction("get_const_vars");
     std::vector<std::string> symbol_const_vars;
@@ -237,10 +371,12 @@ runtime::Module CreateMetadataModule(
   }
 
   if (is_targeting_crt) {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateMetadataModule CreateCrtMetadataModule" << std::endl;
     return CreateCrtMetadataModule(target_module, target, runtime, executor, metadata,
                                    non_crt_exportable_modules, crt_exportable_modules,
                                    const_var_ndarray);
   } else {
+    std::cout << "[ZZDEBUG] tvm::codegen::CreateMetadataModule CreateCppMetadataModule" << std::endl;
     return CreateCppMetadataModule(target_module, target, runtime, metadata, const_vars_by_symbol,
                                    non_crt_exportable_modules, crt_exportable_modules,
                                    const_var_ndarray);
diff --git a/src/tir/usmp/transform/assign_pool_info.cc b/src/tir/usmp/transform/assign_pool_info.cc
index 3acceab6e..272223e95 100644
--- a/src/tir/usmp/transform/assign_pool_info.cc
+++ b/src/tir/usmp/transform/assign_pool_info.cc
@@ -40,22 +40,27 @@ namespace usmp {
  * access.*/
 class PoolInfoAssigner : public StmtExprMutator {
  public:
-  explicit PoolInfoAssigner(const IRModule& module) {
+  // explicit PoolInfoAssigner(const IRModule& module) {
+  // zz: for obfuscated function, used in AssignPoolInfoWithMainFuncName
+  explicit PoolInfoAssigner(const IRModule& module, String main_func_name = ::tvm::runtime::symbol::tvm_module_main) {
     PrimFunc main_func =
-        Downcast<PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
+        // Downcast<PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
+        Downcast<PrimFunc>(module->Lookup(main_func_name)); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
     ICHECK(main_func.defined()) << "main function is not in the module";
     Optional<Target> target_host = main_func->GetAttr<Target>(tvm::attr::kTarget);
     ICHECK(target_host) << "main function does not have a target attr";
     WorkspaceMemoryPools workspace_pools =
         module->GetAttr<WorkspaceMemoryPools>(tvm::attr::kWorkspaceMemoryPools)
-            .value_or(WorkspaceMemoryPools({CreateDefaultWorkspaceMemoryPool(module)}));
+            // .value_or(WorkspaceMemoryPools({CreateDefaultWorkspaceMemoryPool(module)}));
+            .value_or(WorkspaceMemoryPools({CreateDefaultWorkspaceMemoryPool(module, main_func_name)})); // zz: add main_func_name
     // make default ConstantPoolInfo if no constant and no workspace pool infos supplied
     ConstantMemoryPools constant_pools =
         module->GetAttr<ConstantMemoryPools>(tvm::attr::kConstantMemoryPools)
             .value_or(
                 module->GetAttr<WorkspaceMemoryPools>(tvm::attr::kWorkspaceMemoryPools).defined()
                     ? ConstantMemoryPools()
-                    : ConstantMemoryPools({CreateDefaultConstantMemoryPool(module)}));
+                    // : ConstantMemoryPools({CreateDefaultConstantMemoryPool(module)}));
+                    : ConstantMemoryPools({CreateDefaultConstantMemoryPool(module , main_func_name)})); // zz: add main_func_name
     auto to_map = [](auto pool_infos) {
       Map<String, Array<PoolInfo>> pool_map;
       for (const PoolInfo& pool_info : pool_infos) {
@@ -72,8 +77,21 @@ class PoolInfoAssigner : public StmtExprMutator {
     };
 
     target_pool_infos_ = to_map(workspace_pools->pools);
+    // std::cout << "[ZZDEBUG] PoolInfoAssigner::PoolInfoAssigne workspace_pools->pools.size()= (should aways be 1 in unified memory planner) " << workspace_pools->pools.size() << std::endl;
     if (constant_pools.defined()) {
       target_const_pool_infos_ = to_map(constant_pools->pools);
+      // std::cout << "[ZZDEBUG] PoolInfoAssigner::PoolInfoAssigner constant_pools->pools.size()= (should aways be 1 in unified memory planner) " << constant_pools->pools.size() << std::endl;
+      // for (const auto& kv : target_const_pool_infos_) {
+      //   std::cout << "  (zz) PoolInfoAssigner::PoolInfoAssigner target_const_pool_infos_ kv.first= " << kv.first << std::endl;
+      //   for (const auto& pool_info : kv.second) {
+      //     // cast to ConstantPoolInfo
+      //     auto cpool_info = Downcast<ConstantPoolInfo>(pool_info);
+      //     std::cout << "  (zz) PoolInfoAssigner::PoolInfoAssigner target_const_pool_infos_ cpool_info->name= " << cpool_info->pool_name << std::endl;
+      //     // zz: Array<ConstantInfo> constant_info_array is just initialized to [] here
+      //     // zz: constant_info_array is basically parameter (ConstantNode) assigment
+      //     std::cout << "  (zz) PoolInfoAssigner::PoolInfoAssigner target_const_pool_infos_ cpool_info->constant_info_array.size()= " << cpool_info->constant_info_array.size() << std::endl;
+      //   }
+      // }
     }
     mod_ = module->ShallowCopy();
   }
@@ -88,9 +106,14 @@ class PoolInfoAssigner : public StmtExprMutator {
   Map<String, Array<PoolInfo>> target_pool_infos_;
   Map<String, Array<PoolInfo>> target_const_pool_infos_;
   PrimFunc func_;
-  WorkspacePoolInfo CreateDefaultWorkspaceMemoryPool(const IRModule& module);
-  ConstantPoolInfo CreateDefaultConstantMemoryPool(const IRModule& module) {
-    auto p = CreateDefaultWorkspaceMemoryPool(module);
+  // WorkspacePoolInfo CreateDefaultWorkspaceMemoryPool(const IRModule& module);
+  // zz: for obfuscated function, used in PoolInfoAssigner
+  WorkspacePoolInfo CreateDefaultWorkspaceMemoryPool(const IRModule& module, String main_func_name = ::tvm::runtime::symbol::tvm_module_main);
+  // ConstantPoolInfo CreateDefaultConstantMemoryPool(const IRModule& module) {
+  // zz: for obfuscated function, used in PoolInfoAssigner
+  ConstantPoolInfo CreateDefaultConstantMemoryPool(const IRModule& module, String main_func_name = ::tvm::runtime::symbol::tvm_module_main) {
+    // auto p = CreateDefaultWorkspaceMemoryPool(module, main_func_name);
+    auto p = CreateDefaultWorkspaceMemoryPool(module, main_func_name); // zz: add main_func_name
     return ConstantPoolInfo(
         "global_const_workspace", {p->targets}, {},
         PoolInfoProperties(kUnrestrictedPoolSizeHint, kUnknownClockFrequency, kUnknownReadBandwidth,
@@ -98,11 +121,14 @@ class PoolInfoAssigner : public StmtExprMutator {
   }
 };
 
-WorkspacePoolInfo PoolInfoAssigner::CreateDefaultWorkspaceMemoryPool(const tvm::IRModule& module) {
+// WorkspacePoolInfo PoolInfoAssigner::CreateDefaultWorkspaceMemoryPool(const tvm::IRModule& module) {
+// zz: for obfuscated function, used in PoolInfoAssigner
+WorkspacePoolInfo PoolInfoAssigner::CreateDefaultWorkspaceMemoryPool(const tvm::IRModule& module, String main_func_name) {
   VLOG(1) << "Creating default memory pool for:" << std::endl << module;
   Map<Target, String> target_access;
   tir::PrimFunc tir_main_func =
-      Downcast<tir::PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
+      // Downcast<tir::PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
+      Downcast<tir::PrimFunc>(module->Lookup(main_func_name)); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
   Target target_host = tir_main_func->GetAttr<Target>(tvm::attr::kTarget).value();
   for (const auto& kv : module->functions) {
     BaseFunc func = kv.second;
@@ -183,6 +209,14 @@ tvm::transform::Pass AssignPoolInfo() {
   return tvm::transform::CreateModulePass(pass_func, 0, "tir.usmp.AssignPoolInfo", {});
 }
 
+// zz: for obfuscated function, used in PlanMemoryWithMainFuncName
+tvm::transform::Pass AssignPoolInfoWithMainFuncName(String main_func_name) {
+  auto pass_func = [=](IRModule m, tvm::transform::PassContext ctx) {
+    return PoolInfoAssigner(m, main_func_name)();
+  };
+  return tvm::transform::CreateModulePass(pass_func, 0, "tir.usmp.AssignPoolInfo", {});
+}
+
 TVM_REGISTER_GLOBAL("tir.usmp.transform.AssignPoolInfo").set_body_typed(AssignPoolInfo);
 
 }  // namespace transform
diff --git a/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc b/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc
index 0426c0cb1..a83d71b4e 100644
--- a/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc
+++ b/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc
@@ -50,8 +50,11 @@ class PoolAllocationToOffsetConverter : public StmtExprMutator {
  public:
   PoolAllocationToOffsetConverter(const IRModule& module,
                                   const Map<tir::Stmt, PoolAllocation>& pool_allocations,
-                                  bool emit_tvmscript_printable = false)
-      : pool_allocations_(pool_allocations), emit_tvmscript_printable_(emit_tvmscript_printable) {
+                                  // bool emit_tvmscript_printable = false)
+                                  // zz: for obfuscated function, used in ConvertPoolAllocationsToOffsetsWithMainFuncName
+                                  bool emit_tvmscript_printable = false, String main_func_name = ::tvm::runtime::symbol::tvm_module_main)
+      // : pool_allocations_(pool_allocations), emit_tvmscript_printable_(emit_tvmscript_printable) {
+      : pool_allocations_(pool_allocations), emit_tvmscript_printable_(emit_tvmscript_printable), main_func_name_(main_func_name) {
     module_ = module->ShallowCopy();
     for (const auto& kv : pool_allocations) {
       size_t extent_size = -1;
@@ -181,6 +184,9 @@ class PoolAllocationToOffsetConverter : public StmtExprMutator {
 
   Map<PoolInfo, Array<ConstantInfo>> pool_initializations_;
   void AppdendConstInitializationData(ScopeInfo si);
+
+  // zz: for obfuscated function, used in PlanMemoryWithMainFuncName
+  String main_func_name_;
 };
 
 Optional<Var> PoolAllocationToOffsetConverter::GetResourceHandle(const PrimFunc& func) {
@@ -456,7 +462,8 @@ void PoolAllocationToOffsetConverter::AppdendConstInitializationData(
 }
 
 IRModule PoolAllocationToOffsetConverter::operator()() {
-  GlobalVar gv = module_->GetGlobalVar(::tvm::runtime::symbol::tvm_module_main);
+  // GlobalVar gv = module_->GetGlobalVar(::tvm::runtime::symbol::tvm_module_main);
+  GlobalVar gv = module_->GetGlobalVar(main_func_name_); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
   PrimFunc main_func = Downcast<PrimFunc>(module_->Lookup(gv));
   ScopeInfo si = UpdateFunctionScopeInfo(main_func);
   this->scope_stack.push(si);
@@ -499,6 +506,17 @@ tvm::transform::Pass ConvertPoolAllocationsToOffsets(
 TVM_REGISTER_GLOBAL("tir.usmp.transform.ConvertPoolAllocationsToOffsets")
     .set_body_typed(ConvertPoolAllocationsToOffsets);
 
+// zz: for obfuscated function, used in PlanMemoryWithMainFuncName
+tvm::transform::Pass ConvertPoolAllocationsToOffsetsWithMainFuncName(
+    const Map<tir::Stmt, PoolAllocation>& pool_allocations, String main_func_name, Bool emit_tvmscript_printable) {
+  auto pass_func = [=](IRModule m, tvm::transform::PassContext ctx) {
+    return Downcast<IRModule>(PoolAllocationToOffsetConverter(
+        m, pool_allocations, emit_tvmscript_printable->value != 0, main_func_name)());
+  };
+  return tvm::transform::CreateModulePass(pass_func, 0, "tir.usmp.ConvertPoolAllocationsToOffsets",
+                                          {});
+}
+
 }  // namespace transform
 
 }  // namespace usmp
diff --git a/src/tir/usmp/unified_static_memory_planner.cc b/src/tir/usmp/unified_static_memory_planner.cc
index 60030c159..bde33a0a6 100644
--- a/src/tir/usmp/unified_static_memory_planner.cc
+++ b/src/tir/usmp/unified_static_memory_planner.cc
@@ -33,6 +33,8 @@
 #include <tvm/tir/usmp/transform.h>
 #include <tvm/tir/usmp/utils.h>
 
+#include <tvm/relay/base.h> // zz: for PrettyPrint
+
 #include <algorithm>
 #include <string>
 
@@ -82,9 +84,32 @@ IRModule PlanMemory(const IRModule& mod, String algo, bool use_workspace_io,
   Map<BufferInfo, PoolAllocation> buffer_info_pool_allocations =
       algorithm(buffer_info_arr, buffer_info_analysis->memory_pressure);
 
+  // for (const auto& kv : buffer_info_pool_allocations) {
+  //   // check whether kv.second is ConstantPoolInfo
+  //   if (kv.second->pool_info->IsInstance<ConstantPoolInfoNode>()) {
+  //     auto cpool_info = Downcast<ConstantPoolInfo>(kv.second->pool_info);
+  //     std::cout << "[ZZDEBUG] PlanMemory buffer_info_pool_allocations cpool_info->name= " << cpool_info->pool_name << std::endl;
+  //     std::cout << "[ZZDEBUG] PlanMemory buffer_info_pool_allocations cpool_info->constant_info_array.size()= (here is still 0)" << cpool_info->constant_info_array.size() << std::endl;
+
+  //     std::cout << "[ZZDEBUG] PlanMemory buffer_info_pool_allocations kv.second->pool_info->byte_offset= (offset is determined)" << kv.second->byte_offset->value << std::endl;
+  //   }
+  // }
+
   Map<Stmt, PoolAllocation> stmt_pool_allocations = AssignStmtPoolAllocations(
       buffer_info_analysis->buffer_info_stmts, buffer_info_pool_allocations);
 
+  // // zz: sam as buffer_info_pool_allocations
+  // for (const auto& kv : stmt_pool_allocations) {
+  //   // check whether kv.second is ConstantPoolInfo
+  //   if (kv.second->pool_info->IsInstance<ConstantPoolInfoNode>()) {
+  //     auto cpool_info = Downcast<ConstantPoolInfo>(kv.second->pool_info);
+  //     std::cout << "[ZZDEBUG] PlanMemory stmt_pool_allocations cpool_info->name= " << cpool_info->pool_name << std::endl;
+  //     std::cout << "[ZZDEBUG] PlanMemory stmt_pool_allocations cpool_info->constant_info_array.size()= (same as above)" << cpool_info->constant_info_array.size() << std::endl;
+
+  //     std::cout << "[ZZDEBUG] PlanMemory stmt_pool_allocations kv.second->pool_info->byte_offset= (same as above)" << kv.second->byte_offset->value << std::endl;
+  //   }
+  // }
+
   module = transform::ConvertPoolAllocationsToOffsets(stmt_pool_allocations)(module);
   if (use_workspace_io) {
     Map<String, PoolAllocation> io_pool_allocations =
@@ -103,6 +128,62 @@ IRModule PlanMemory(const IRModule& mod, String algo, bool use_workspace_io,
   return module;
 }
 
+// zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+IRModule PlanMemoryWithMainFuncName(const IRModule& mod, String algo, bool use_workspace_io,
+                    Optional<String> opt_custom_algo, String main_func_name) {
+  VLOG(1) << "workspace required = " << CalculateModuleWorkspaceSizeWithMainFuncName(mod, main_func_name); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
+
+  IRModule module = mod->ShallowCopy();
+  if (use_workspace_io) {
+    module = transform::CreateAllocatesForIO()(module);
+  }
+  module = transform::AssignPoolInfoWithMainFuncName(main_func_name)(module); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
+  PrimFunc main_func = Downcast<PrimFunc>(module->Lookup(main_func_name)); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
+  BufferInfoAnalysis buffer_info_analysis = ExtractBufferInfo(main_func, module);
+  // std::cout << "[ZZDEBUG] PlanMemoryWithMainFuncName after ExtractBufferInfo" << std::endl;
+
+  Array<BufferInfo> buffer_info_arr =
+      ConvertToArrayOfBufferInfo(buffer_info_analysis->buffer_info_stmts);
+  // std::cout << "[ZZDEBUG] PlanMemoryWithMainFuncName after ConvertToArrayOfBufferInfo" << std::endl;
+
+  decltype(algorithms)::mapped_type algorithm;
+  if (opt_custom_algo) {
+    String algo_func_name = "tir.usmp.algo." + opt_custom_algo.value();
+    const runtime::PackedFunc* pfAlgo = runtime::Registry::Get(algo_func_name);
+    CHECK(pfAlgo) << "The selected custom USMP algorithm : " << opt_custom_algo.value()
+                  << " is not defined. Please register it as " << algo_func_name;
+    algorithm = *pfAlgo;
+  } else {
+    CHECK(algorithms.count(algo))
+        << "The selected USMP algorithm : " << algo
+        << " is not defined. Please define it in the above algorithms map.";
+    algorithm = algorithms[algo];
+  }
+  Map<BufferInfo, PoolAllocation> buffer_info_pool_allocations =
+      algorithm(buffer_info_arr, buffer_info_analysis->memory_pressure);
+
+  Map<Stmt, PoolAllocation> stmt_pool_allocations = AssignStmtPoolAllocations(
+      buffer_info_analysis->buffer_info_stmts, buffer_info_pool_allocations);
+
+  module = transform::ConvertPoolAllocationsToOffsetsWithMainFuncName(stmt_pool_allocations, main_func_name)(module); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
+  if (use_workspace_io) {
+    Map<String, PoolAllocation> io_pool_allocations =
+        GetIOPoolAllocations(buffer_info_pool_allocations);
+    module = WithAttr(module, tvm::attr::kIOTensorPoolAllocations, io_pool_allocations);
+  }
+  tir::PrimFunc tir_main_func =
+      Downcast<tir::PrimFunc>(module->Lookup(main_func_name)); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
+
+  Optional<Array<tir::usmp::AllocatedPoolInfo>> allocated_pool_infos =
+      tir_main_func->GetAttr<Array<tir::usmp::AllocatedPoolInfo>>(tvm::attr::kPoolArgs);
+  if (allocated_pool_infos) {
+    for (const tir::usmp::AllocatedPoolInfo& allocated_pool_info : allocated_pool_infos.value()) {
+      VLOG(1) << "pool_size = " << allocated_pool_info->allocated_size;
+    }
+  }
+  return module;
+}
+
 }  // namespace usmp
 
 namespace transform {
@@ -135,6 +216,33 @@ tvm::transform::Pass UnifiedStaticMemoryPlanner() {
 TVM_REGISTER_GLOBAL("tir.transform.UnifiedStaticMemoryPlanner")
     .set_body_typed(UnifiedStaticMemoryPlanner);
 
+// zz: for obfuscated function, used in AOTExecutorCodegen::CodegenObfCall
+tvm::transform::Pass UnifiedStaticMemoryPlannerWithMainFuncName(String main_func_name) {
+  auto usmp_main_pass_func = [=](IRModule m, tvm::transform::PassContext ctx) {
+    auto algorithm_str = ctx->GetConfig(kUSMPAlgorithmOption, String(usmp::kDefaultAlgo));
+    auto use_workspace_io = ctx->GetConfig(kUSMPUseWorkspaceIO, Bool(false));
+    auto custom_algorithm_str = ctx->GetConfig<String>(kUSMPCustomAlgorithmOption);
+    tvm::relay::Executor executor_config =
+        m->GetAttr<tvm::relay::Executor>(tvm::attr::kExecutor).value();
+    String interface_api = executor_config->GetAttr<String>("interface-api").value_or("packed");
+    tvm::relay::Runtime runtime_config =
+        m->GetAttr<tvm::relay::Runtime>(tvm::attr::kRuntime).value();
+    if (use_workspace_io.value()) {
+      CHECK(interface_api == "c") << kUSMPUseWorkspaceIO
+                                  << " option is only compatible with interface_api c.\n"
+                                  << "Please use interface_api c to be able to enable "
+                                  << kUSMPUseWorkspaceIO << "\n";
+    }
+    return Downcast<IRModule>(
+        usmp::PlanMemoryWithMainFuncName(m, algorithm_str.value_or(String(usmp::kDefaultAlgo)),
+                         use_workspace_io.value_or(Bool(false)), custom_algorithm_str, main_func_name));
+  };
+
+  return tvm::transform::CreateModulePass(usmp_main_pass_func, 0,
+                                          "tir.transform.UnifiedStaticMemoryPlanner", {});
+}
+
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm
diff --git a/src/tir/usmp/utils.cc b/src/tir/usmp/utils.cc
index 88a649685..f6d1d66c3 100644
--- a/src/tir/usmp/utils.cc
+++ b/src/tir/usmp/utils.cc
@@ -217,6 +217,20 @@ class ModuleWorkspaceSizeCalculator : public StmtExprVisitor {
     target_host_ = target_host.value();
   }
 
+  // zz: for obfuscated function, used in CalculateModuleWorkspaceSize
+  explicit ModuleWorkspaceSizeCalculator(const IRModule& module, String main_func_name) : mod_(module) {
+    for (const auto& gv_func : mod_->functions) {
+      if ((gv_func.second)->IsInstance<tir::PrimFuncNode>()) {
+        functions_.Set(gv_func.first->name_hint, Downcast<PrimFunc>(gv_func.second));
+      }
+    }
+    main_func_ = Downcast<PrimFunc>(module->Lookup(main_func_name)); // zz: [originally operate on ::tvm::runtime::symbol::tvm_module_main]
+    ICHECK(main_func_.defined()) << "main function is not in the module";
+    Optional<Target> target_host = main_func_->GetAttr<Target>(tvm::attr::kTarget);
+    ICHECK(target_host) << "main function does not have a target attr";
+    target_host_ = target_host.value();
+  }
+
   Integer operator()() {
     UpdateWorkspaceData(main_func_);
     return Integer(max_workspace_size);
@@ -264,6 +278,11 @@ Integer CalculateModuleWorkspaceSize(const IRModule& mod) {
   return ModuleWorkspaceSizeCalculator(mod)();
 }
 
+// zz: for obfuscated function, used in PlanMemoryWithMainFuncName
+Integer CalculateModuleWorkspaceSizeWithMainFuncName(const IRModule& mod, String main_func_name) {
+  return ModuleWorkspaceSizeCalculator(mod, main_func_name)();
+}
+
 TVM_REGISTER_GLOBAL("tir.usmp.CreateArrayBufferInfo")
     .set_body_typed([](Map<BufferInfo, Stmt> buffer_info_map) {
       return (ConvertToArrayOfBufferInfo(buffer_info_map));