From 23cf14c4b59fc8dfc8d696ee50cd86a6122f0686 Mon Sep 17 00:00:00 2001
From: Mingchuan Wu <wumingchuan1992@foxmail.com>
Date: Wed, 1 Mar 2023 16:04:51 +0800
Subject: [PATCH 17/23] [Refactoring] Modify API parameters to improve
 usability.


diff --git a/include/Dialect/PluginOps.td b/include/Dialect/PluginOps.td
index 93ba116..47909c7 100644
--- a/include/Dialect/PluginOps.td
+++ b/include/Dialect/PluginOps.td
@@ -95,7 +95,7 @@ def LoopOp : Plugin_Op<"loop", [NoSideEffect]> {
     bool IsLoopFather(mlir::Block*);
     std::vector<std::pair<mlir::Block*, mlir::Block*> > GetExitEdges();
     std::vector<mlir::Block*> GetLoopBody(); 
-    void AddLoop(uint64_t, uint64_t);
+    void AddLoop(LoopOp *, FunctionOp *);
     void AddBlock(mlir::Block*);
   }];
 }
@@ -118,9 +118,8 @@ def CallOp : Plugin_Op<"call", [
     OpBuilderDAG<(ins "int64_t":$id,
                       "ArrayRef<Value>":$arguments)>,
     // Only for server.
-    OpBuilderDAG<(ins "Value":$func,
-                      "ArrayRef<Value>":$arguments, "Block *":$block)>,
-    OpBuilderDAG<(ins "ArrayRef<Value>":$arguments, "Block *":$block)>
+    OpBuilderDAG<(ins "Value":$func, "ArrayRef<Value>":$arguments)>,
+    OpBuilderDAG<(ins "ArrayRef<Value>":$arguments)>
   ];
   let extraClassDeclaration = [{
     bool SetLHS(Value lhs);
@@ -160,7 +159,7 @@ def AssignOp : Plugin_Op<"assign", [NoSideEffect]> {
                       "uint64_t":$id, "IExprCode":$exprCode)>,
     // Only for server.
     OpBuilderDAG<(ins "ArrayRef<Value>":$operands,
-                      "IExprCode":$exprCode, "Block *":$block)>
+                      "IExprCode":$exprCode)>
   ];
   let extraClassDeclaration = [{
     Value GetLHS() { return getOperand(0); }
@@ -625,7 +624,7 @@ def FallThroughOp : Plugin_TerminatorOp<"fallthrough", [NoSideEffect]> {
   let builders = [
     OpBuilderDAG<(ins "uint64_t":$address, "Block*":$dest, "uint64_t":$destaddr)>,
     // Only for server.
-    OpBuilderDAG<(ins "uint64_t":$address, "Block*":$dest)>
+    OpBuilderDAG<(ins "Block*":$src, "Block*":$dest)>
   ];
 }
 
@@ -646,8 +645,8 @@ def CondOp : Plugin_TerminatorOp<"condition", [NoSideEffect]> {
                       "uint64_t":$tbaddr, "uint64_t":$fbaddr, "Value":$trueLabel,
                       "Value":$falseLabel)>,
     // Only for server.
-    OpBuilderDAG<(ins "IComparisonCode":$condCode, "Value":$lhs, "Value":$rhs, 
-                      "Block*":$tb, "Block*":$fb, "Block *":$block)>
+    OpBuilderDAG<(ins "IComparisonCode":$condCode, "Value":$lhs,
+                      "Value":$rhs, "Block*":$tb, "Block*":$fb)>
   ];
   let extraClassDeclaration = [{
     Value GetLHS() { return getOperand(0); }
diff --git a/include/PluginAPI/BasicPluginOpsAPI.h b/include/PluginAPI/BasicPluginOpsAPI.h
index 923dd44..6f0a4af 100644
--- a/include/PluginAPI/BasicPluginOpsAPI.h
+++ b/include/PluginAPI/BasicPluginOpsAPI.h
@@ -52,13 +52,13 @@ public:
     virtual vector<LoopOp> GetLoopsFromFunc(uint64_t) = 0;
     virtual LoopOp GetLoopById(uint64_t) = 0;
     virtual void AddLoop(uint64_t, uint64_t, uint64_t) = 0;
-    virtual void AddBlockToLoop(uint64_t, uint64_t) = 0;
+    virtual void AddBlockToLoop(mlir::Block*, LoopOp *) = 0;
     virtual void DeleteLoop(uint64_t) = 0;
     virtual vector<mlir::Block*> GetLoopBody(uint64_t) = 0;
     virtual bool IsBlockInLoop(uint64_t, uint64_t) = 0;
     virtual pair<mlir::Block*, mlir::Block*> LoopSingleExit(uint64_t) = 0;
     virtual vector<pair<mlir::Block*, mlir::Block*> > GetLoopExitEdges(uint64_t) = 0;
-    virtual LoopOp GetBlockLoopFather(uint64_t) = 0;
+    virtual LoopOp GetBlockLoopFather(mlir::Block*) = 0;
     virtual PhiOp GetPhiOp(uint64_t) = 0;
     virtual CallOp GetCallOp(uint64_t) = 0;
     virtual bool SetLhsInCallOp(uint64_t, uint64_t) = 0;
@@ -80,7 +80,7 @@ public:
     virtual mlir::Value CreateSSAOp(mlir::Type) = 0;
     virtual mlir::Value ConfirmValue(mlir::Value) = 0;
     virtual mlir::Value BuildMemRef(PluginIR::PluginTypeBase, mlir::Value, mlir::Value) = 0;
-    virtual bool RedirectFallthroughTarget(FallThroughOp&, uint64_t, uint64_t) = 0;
+    virtual bool RedirectFallthroughTarget(FallThroughOp&, mlir::Block*, mlir::Block*) = 0;
     virtual mlir::Operation* GetSSADefOperation(uint64_t) = 0;
 }; // class BasicPluginOpsAPI
 } // namespace PluginAPI
diff --git a/include/PluginAPI/ControlFlowAPI.h b/include/PluginAPI/ControlFlowAPI.h
index a436e85..429ba11 100644
--- a/include/PluginAPI/ControlFlowAPI.h
+++ b/include/PluginAPI/ControlFlowAPI.h
@@ -41,15 +41,16 @@ public:
     bool UpdateSSA(void);
     vector<PhiOp> GetAllPhiOpInsideBlock(mlir::Block *b);
 
-    uint64_t CreateBlock(mlir::Block*, uint64_t, uint64_t);
-    void DeleteBlock(mlir::Block*, uint64_t, uint64_t);
-    void SetImmediateDominator(uint64_t, uint64_t, uint64_t);
+    mlir::Block* CreateBlock(mlir::Block*, FunctionOp*);
+    void DeleteBlock(mlir::Block*, FunctionOp*);
+    void SetImmediateDominator(uint64_t, mlir::Block*, uint64_t);
+    void SetImmediateDominator(uint64_t, mlir::Block*, mlir::Block*);
     uint64_t GetImmediateDominator(uint64_t, uint64_t);
-    uint64_t RecomputeDominator(uint64_t, uint64_t);
+    uint64_t RecomputeDominator(uint64_t, mlir::Block*);
     mlir::Value CreateNewDef(mlir::Value, mlir::Operation *);
 
     void CreateFallthroughOp(uint64_t, uint64_t);
-    void RemoveEdge(uint64_t, uint64_t);
+    void RemoveEdge(mlir::Block*, mlir::Block*);
 private:
     bool GetUpdateOperationResult(const string &funName);
     vector<PhiOp> GetPhiOperationResult(const string &funName, const string& params);
diff --git a/include/PluginAPI/PluginServerAPI.h b/include/PluginAPI/PluginServerAPI.h
index 3da28f5..0d6a5bc 100644
--- a/include/PluginAPI/PluginServerAPI.h
+++ b/include/PluginAPI/PluginServerAPI.h
@@ -48,15 +48,15 @@ public:
     bool IsBlockInLoop(uint64_t loopID, uint64_t blockID);
     void DeleteLoop(uint64_t loopID);
     void AddLoop(uint64_t loopID, uint64_t outerID, uint64_t funcID);
-    void AddBlockToLoop(uint64_t blockID, uint64_t loopID);
+    void AddBlockToLoop(mlir::Block*, LoopOp *);
     pair<mlir::Block*, mlir::Block*> LoopSingleExit(uint64_t loopID);
     vector<pair<mlir::Block*, mlir::Block*> > GetLoopExitEdges(uint64_t loopID);
     mlir::Block* GetHeader(uint64_t loopID);
     mlir::Block* GetLatch(uint64_t loopID);
-    void SetHeader(uint64_t loopID, uint64_t blockID);
-    void SetLatch(uint64_t loopID, uint64_t blockID);
+    void SetHeader(LoopOp*, mlir::Block*);
+    void SetLatch(LoopOp*, mlir::Block*);
     vector<mlir::Block*> GetLoopBody(uint64_t loopID);
-    LoopOp GetBlockLoopFather(uint64_t blockID);
+    LoopOp GetBlockLoopFather(mlir::Block*);
     mlir::Block* FindBlock(uint64_t);
     uint64_t FindBasicBlock(mlir::Block*);
     bool InsertValue(uint64_t, mlir::Value);
@@ -84,7 +84,7 @@ public:
     mlir::Value CreateSSAOp(mlir::Type) override;
     mlir::Value ConfirmValue(mlir::Value);
     mlir::Value BuildMemRef(PluginIR::PluginTypeBase, mlir::Value, mlir::Value);
-    bool RedirectFallthroughTarget(FallThroughOp&, uint64_t, uint64_t) override;
+    bool RedirectFallthroughTarget(FallThroughOp&, mlir::Block*, mlir::Block*) override;
     mlir::Operation* GetSSADefOperation(uint64_t) override;
     void InsertCreatedBlock(uint64_t id, mlir::Block* block);
 
diff --git a/lib/Dialect/PluginOps.cpp b/lib/Dialect/PluginOps.cpp
index ab0e3ce..5beb558 100644
--- a/lib/Dialect/PluginOps.cpp
+++ b/lib/Dialect/PluginOps.cpp
@@ -163,17 +163,13 @@ Block* LoopOp::GetLatch()
 void LoopOp::SetHeader(mlir::Block* b)
 {
     PluginAPI::PluginServerAPI pluginAPI;
-    uint64_t loopId = idAttr().getInt();
-    uint64_t blockId = pluginAPI.FindBasicBlock(b);
-    pluginAPI.SetHeader(loopId, blockId);
+    pluginAPI.SetHeader(this, b);
 }
 
 void LoopOp::SetLatch(mlir::Block* b)
 {
     PluginAPI::PluginServerAPI pluginAPI;
-    uint64_t loopId = idAttr().getInt();
-    uint64_t blockId = pluginAPI.FindBasicBlock(b);
-    pluginAPI.SetLatch(loopId, blockId);
+    pluginAPI.SetLatch(this, b);
 }
 
 vector<mlir::Block*> LoopOp::GetLoopBody()
@@ -223,8 +219,7 @@ bool LoopOp::IsLoopFather(mlir::Block* b)
 {
     PluginAPI::PluginServerAPI pluginAPI;
     uint64_t loopId = idAttr().getInt();
-    uint64_t blockId = pluginAPI.FindBasicBlock(b);
-    LoopOp loopFather = pluginAPI.GetBlockLoopFather(blockId);
+    LoopOp loopFather = pluginAPI.GetBlockLoopFather(b);
     uint64_t id = loopFather.idAttr().getInt();
     return id == loopId;
 }
@@ -236,19 +231,18 @@ vector<pair<mlir::Block*, mlir::Block*> > LoopOp::GetExitEdges()
     return pluginAPI.GetLoopExitEdges(loopId);
 }
 
-void LoopOp::AddLoop(uint64_t outerId, uint64_t funcId)
+void LoopOp::AddLoop(LoopOp* outerLoop, FunctionOp* funcOp)
 {
     PluginAPI::PluginServerAPI pluginAPI;
     uint64_t loopId = idAttr().getInt();
-    return pluginAPI.AddLoop(loopId, outerId, funcId);
+    return pluginAPI.AddLoop(loopId, outerLoop->idAttr().getInt(),
+                             funcOp->idAttr().getInt());
 }
 
 void LoopOp::AddBlock(mlir::Block* block)
 {
     PluginAPI::PluginServerAPI pluginAPI;
-    uint64_t blockId = pluginAPI.FindBasicBlock(block);
-    uint64_t loopId = idAttr().getInt();
-    pluginAPI.AddBlockToLoop(blockId, loopId);
+    pluginAPI.AddBlockToLoop(block, this);
 }
 
 // ===----------------------------------------------------------------------===//
@@ -595,9 +589,12 @@ bool CallOp::SetLHS(Value lhs)
 }
 
 void CallOp::build(OpBuilder &builder, OperationState &state,
-                   Value func, ArrayRef<Value> arguments, Block * block)
+                   Value func, ArrayRef<Value> arguments)
 {
+    Block *insertionBlock = builder.getInsertionBlock();
+    assert(insertionBlock && "No InsertPoint is set for the OpBuilder.");
     PluginAPI::PluginServerAPI pluginAPI;
+    uint64_t blockId = pluginAPI.FindBasicBlock(insertionBlock);
     PlaceholderOp funcOp = func.getDefiningOp<PlaceholderOp>();
     uint64_t funcId = funcOp.idAttr().getInt();
     vector<uint64_t> argIds;
@@ -605,7 +602,6 @@ void CallOp::build(OpBuilder &builder, OperationState &state,
         uint64_t argId = GetValueId(v);
         argIds.push_back(argId);
     }
-    uint64_t blockId = pluginAPI.FindBasicBlock(block);
     uint64_t id = pluginAPI.CreateCallOp(blockId, funcId, argIds);
     state.addAttribute("id", builder.getI64IntegerAttr(id));
     state.addOperands(arguments);
@@ -614,15 +610,17 @@ void CallOp::build(OpBuilder &builder, OperationState &state,
 }
 
 void CallOp::build(OpBuilder &builder, OperationState &state,
-                   ArrayRef<Value> arguments, Block * block)
+                   ArrayRef<Value> arguments)
 {
+    Block *insertionBlock = builder.getInsertionBlock();
+    assert(insertionBlock && "No InsertPoint is set for the OpBuilder.");
     PluginAPI::PluginServerAPI pluginAPI;
+    uint64_t blockId = pluginAPI.FindBasicBlock(insertionBlock);
     vector<uint64_t> argIds;
     for (auto v : arguments) {
         uint64_t argId = GetValueId(v);
         argIds.push_back(argId);
     }
-    uint64_t blockId = pluginAPI.FindBasicBlock(block);
     uint64_t funcId = 0;
     uint64_t id = pluginAPI.CreateCallOp(blockId, funcId, argIds);
     state.addAttribute("id", builder.getI64IntegerAttr(id));
@@ -653,12 +651,14 @@ void CondOp::build(OpBuilder &builder, OperationState &state,
 
 void CondOp::build(OpBuilder &builder, OperationState &state,
                    IComparisonCode condCode, Value lhs, Value rhs, Block* tb,
-                   Block* fb, Block * block)
+                   Block* fb)
 {
+    Block *insertionBlock = builder.getInsertionBlock();
+    assert(insertionBlock && "No InsertPoint is set for the OpBuilder.");
     PluginAPI::PluginServerAPI pluginAPI;
+    uint64_t blockId = pluginAPI.FindBasicBlock(insertionBlock);
     uint64_t lhsId = GetValueId(lhs);
     uint64_t rhsId = GetValueId(rhs);
-    uint64_t blockId = pluginAPI.FindBasicBlock(block);
     uint64_t tbaddr = pluginAPI.FindBasicBlock(tb);
     uint64_t fbaddr = pluginAPI.FindBasicBlock(fb);
     uint64_t id = pluginAPI.CreateCondOp(blockId, condCode, lhsId, rhsId,
@@ -737,16 +737,17 @@ void AssignOp::build(OpBuilder &builder, OperationState &state,
 }
 
 void AssignOp::build(OpBuilder &builder, OperationState &state,
-                     ArrayRef<Value> operands, IExprCode exprCode,
-                     Block * block)
+                     ArrayRef<Value> operands, IExprCode exprCode)
 {
+    Block *insertionBlock = builder.getInsertionBlock();
+    assert(insertionBlock && "No InsertPoint is set for the OpBuilder.");
     PluginAPI::PluginServerAPI pluginAPI;
+    uint64_t blockId = pluginAPI.FindBasicBlock(insertionBlock);
     vector<uint64_t> argIds;
     for (auto v : operands) {
         uint64_t argId = GetValueId(v);
         argIds.push_back(argId);
     }
-    uint64_t blockId = pluginAPI.FindBasicBlock(block);
     uint64_t id = pluginAPI.CreateAssignOp(blockId, exprCode, argIds);
     state.addAttribute("id", builder.getI64IntegerAttr(id));
     state.addAttribute("exprCode",
@@ -821,7 +822,7 @@ void SwitchOp::build(OpBuilder &builder, OperationState &state,
     state.addAttribute("address", builder.getI64IntegerAttr(address));
     state.addAttribute("defaultaddr", builder.getI64IntegerAttr(defaultaddr));
     llvm::SmallVector<mlir::Attribute, 4> attributes;
-    for (int i = 0; i < caseaddr.size(); ++i) {
+    for (size_t i = 0; i < caseaddr.size(); ++i) {
         attributes.push_back(builder.getI64IntegerAttr(caseaddr[i]));
     }
     state.addAttribute("caseaddrs", builder.getArrayAttr(attributes));
@@ -842,9 +843,10 @@ void FallThroughOp::build(OpBuilder &builder, OperationState &state,
 }
 
 void FallThroughOp::build(OpBuilder &builder, OperationState &state,
-                          uint64_t address, Block* dest)
+                          Block* src, Block* dest)
 {
     PluginAPI::PluginServerAPI pluginAPI;
+    uint64_t address = pluginAPI.FindBasicBlock(src);
     uint64_t destaddr = pluginAPI.FindBasicBlock(dest);
 
     PluginAPI::ControlFlowAPI cfgAPI;
@@ -892,7 +894,7 @@ void TransactionOp::build(OpBuilder &builder, OperationState &state, uint64_t id
     state.addAttribute("id", builder.getI64IntegerAttr(id));
     state.addAttribute("address", builder.getI64IntegerAttr(address));
     llvm::SmallVector<mlir::Attribute, 4> attributes;
-    for (int i = 0; i < stmtaddr.size(); ++i) {
+    for (size_t i = 0; i < stmtaddr.size(); ++i) {
         attributes.push_back(builder.getI64IntegerAttr(stmtaddr[i]));
     }
     state.addAttribute("stmtaddr", builder.getArrayAttr(attributes));
@@ -932,7 +934,7 @@ void EHDispatchOp::build(OpBuilder &builder, OperationState &state, uint64_t id,
     state.addAttribute("region", builder.getI64IntegerAttr(region));
     state.addSuccessors(ehHandlers);
     llvm::SmallVector<mlir::Attribute, 4> attributes;
-    for (int i = 0; i < ehHandlersaddrs.size(); ++i) {
+    for (size_t i = 0; i < ehHandlersaddrs.size(); ++i) {
         attributes.push_back(builder.getI64IntegerAttr(ehHandlersaddrs[i]));
     }
     state.addAttribute("ehHandlersaddrs", builder.getArrayAttr(attributes));
@@ -946,7 +948,7 @@ void BindOp::build(OpBuilder &builder, OperationState &state, uint64_t id, Value
     state.addAttribute("id", builder.getI64IntegerAttr(id));
     state.addOperands({vars, block});
     llvm::SmallVector<mlir::Attribute, 4> attributes;
-    for (int i = 0; i < body.size(); ++i) {
+    for (size_t i = 0; i < body.size(); ++i) {
         attributes.push_back(builder.getI64IntegerAttr(body[i]));
     }
     state.addAttribute("body", builder.getArrayAttr(attributes));
@@ -960,12 +962,12 @@ void TryOp::build(OpBuilder &builder, OperationState &state, uint64_t id, ArrayR
 {
     state.addAttribute("id", builder.getI64IntegerAttr(id));
     llvm::SmallVector<mlir::Attribute, 4> attributes;
-    for (int i = 0; i < eval.size(); ++i) {
+    for (size_t i = 0; i < eval.size(); ++i) {
         attributes.push_back(builder.getI64IntegerAttr(eval[i]));
     }
     state.addAttribute("eval", builder.getArrayAttr(attributes));
     attributes.clear();
-    for (int i = 0; i < cleanup.size(); ++i) {
+    for (size_t i = 0; i < cleanup.size(); ++i) {
         attributes.push_back(builder.getI64IntegerAttr(cleanup[i]));
     }
     state.addAttribute("cleanup", builder.getArrayAttr(attributes));
@@ -980,7 +982,7 @@ void CatchOp::build(OpBuilder &builder, OperationState &state, uint64_t id, Valu
     state.addAttribute("id", builder.getI64IntegerAttr(id));
     state.addOperands(types);
     llvm::SmallVector<mlir::Attribute, 4> attributes;
-    for (int i = 0; i < handler.size(); ++i) {
+    for (size_t i = 0; i < handler.size(); ++i) {
         attributes.push_back(builder.getI64IntegerAttr(handler[i]));
     }
     state.addAttribute("handler", builder.getArrayAttr(attributes));
diff --git a/lib/PluginAPI/ControlFlowAPI.cpp b/lib/PluginAPI/ControlFlowAPI.cpp
index 720bac5..b356ae9 100644
--- a/lib/PluginAPI/ControlFlowAPI.cpp
+++ b/lib/PluginAPI/ControlFlowAPI.cpp
@@ -110,24 +110,30 @@ vector<PhiOp> ControlFlowAPI::GetAllPhiOpInsideBlock(mlir::Block *b)
     return GetPhiOperationResult(funName, params);
 }
 
-uint64_t ControlFlowAPI::CreateBlock(mlir::Block* b, uint64_t funcAddr, uint64_t bbAddr)
+mlir::Block* ControlFlowAPI::CreateBlock(mlir::Block* b, FunctionOp *funcOp)
 {
     Json::Value root;
     string funName = __func__;
+    uint64_t funcAddr = funcOp->idAttr().getInt();
     assert(funcAddr);
+    PluginServer *server = PluginServer::GetInstance();
+    uint64_t bbAddr = server->FindBasicBlock(b);
     assert(bbAddr);
     root["funcaddr"] = std::to_string(funcAddr);
     root["bbaddr"] = std::to_string(bbAddr);
     string params = root.toStyledString();
-    PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
-    return PluginServer::GetInstance()->GetBlockResult(b);
+    server->RemoteCallClientWithAPI(funName, params);
+    uint64_t retId = server->GetBlockResult(b);
+    return server->FindBlock(retId);
 }
 
-void ControlFlowAPI::DeleteBlock(mlir::Block* b, uint64_t funcAddr, uint64_t bbAddr)
+void ControlFlowAPI::DeleteBlock(mlir::Block* b, FunctionOp* funcOp)
 {
     Json::Value root;
     string funName = __func__;
+    uint64_t funcAddr = funcOp->idAttr().getInt();
     assert(funcAddr);
+    uint64_t bbAddr = PluginServer::GetInstance()->FindBasicBlock(b);
     assert(bbAddr);
     root["funcaddr"] = std::to_string(funcAddr);
     root["bbaddr"] = std::to_string(bbAddr);
@@ -137,10 +143,29 @@ void ControlFlowAPI::DeleteBlock(mlir::Block* b, uint64_t funcAddr, uint64_t bbA
 }
 
 /* dir: 1 or 2 */
-void ControlFlowAPI::SetImmediateDominator(uint64_t dir, uint64_t bbAddr, uint64_t domiAddr)
+void ControlFlowAPI::SetImmediateDominator(uint64_t dir, mlir::Block* bb, uint64_t domiAddr)
+{
+    Json::Value root;
+    string funName = __func__;
+    PluginServer *server = PluginServer::GetInstance();
+    uint64_t bbAddr = server->FindBasicBlock(bb);
+    if (!bbAddr || !domiAddr) return;
+    assert(dir && bbAddr && domiAddr);
+    root["dir"] = std::to_string(dir);
+    root["bbaddr"] = std::to_string(bbAddr);
+    root["domiaddr"] = std::to_string(domiAddr);
+    string params = root.toStyledString();
+    PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
+}
+
+void ControlFlowAPI::SetImmediateDominator(uint64_t dir, mlir::Block* bb,
+                                           mlir::Block* domi)
 {
     Json::Value root;
     string funName = __func__;
+    PluginServer *server = PluginServer::GetInstance();
+    uint64_t bbAddr = server->FindBasicBlock(bb);
+    uint64_t domiAddr = server->FindBasicBlock(domi);
     if (!bbAddr || !domiAddr) return;
     assert(dir && bbAddr && domiAddr);
     root["dir"] = std::to_string(dir);
@@ -163,10 +188,11 @@ uint64_t ControlFlowAPI::GetImmediateDominator(uint64_t dir, uint64_t bbAddr)
 }
 
 /* dir: 1 or 2 */
-uint64_t ControlFlowAPI::RecomputeDominator(uint64_t dir, uint64_t bbAddr)
+uint64_t ControlFlowAPI::RecomputeDominator(uint64_t dir, mlir::Block* bb)
 {
     Json::Value root;
     string funName = __func__;
+    uint64_t bbAddr = PluginServer::GetInstance()->FindBasicBlock(bb);
     assert(dir && bbAddr);
     root["dir"] = std::to_string(dir);
     root["bbaddr"] = std::to_string(bbAddr);
@@ -201,10 +227,12 @@ void ControlFlowAPI::CreateFallthroughOp(
     PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
 }
 
-void ControlFlowAPI::RemoveEdge(uint64_t src, uint64_t dest)
+void ControlFlowAPI::RemoveEdge(mlir::Block* srcBB, mlir::Block* destBB)
 {
     Json::Value root;
     string funName = __func__;
+    uint64_t src = PluginServer::GetInstance()->FindBasicBlock(srcBB);
+    uint64_t dest = PluginServer::GetInstance()->FindBasicBlock(destBB);
     root["src"] = std::to_string(src);
     root["dest"] = std::to_string(dest);
     string params = root.toStyledString();
diff --git a/lib/PluginAPI/PluginServerAPI.cpp b/lib/PluginAPI/PluginServerAPI.cpp
index ff9c90e..6a52dd6 100644
--- a/lib/PluginAPI/PluginServerAPI.cpp
+++ b/lib/PluginAPI/PluginServerAPI.cpp
@@ -486,12 +486,12 @@ void PluginServerAPI::AddLoop(uint64_t loopID, uint64_t outerID, uint64_t funcID
     PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
 }
 
-void PluginServerAPI::AddBlockToLoop(uint64_t blockID, uint64_t loopID)
+void PluginServerAPI::AddBlockToLoop(mlir::Block* b, LoopOp *loop)
 {
     Json::Value root;
     string funName("AddBlockToLoop");
-    root["blockId"] = blockID;
-    root["loopId"] = loopID;
+    root["blockId"] = PluginServer::GetInstance()->FindBasicBlock(b);
+    root["loopId"] = loop->idAttr().getInt();
     string params = root.toStyledString();
     PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
 }
@@ -527,22 +527,24 @@ mlir::Block* PluginServerAPI::GetLatch(uint64_t loopID)
     return BlockResult(funName, params);
 }
 
-void PluginServerAPI::SetHeader(uint64_t loopID, uint64_t blockID)
+void PluginServerAPI::SetHeader(LoopOp * loop, mlir::Block* b)
 {
     Json::Value root;
     string funName("SetHeader");
-    root["loopId"] = std::to_string(loopID);
-    root["blockId"] = std::to_string(blockID);
+    root["loopId"] = std::to_string(loop->idAttr().getInt());
+    root["blockId"] = std::to_string(
+        PluginServer::GetInstance()->FindBasicBlock(b));
     string params = root.toStyledString();
     PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
 }
 
-void PluginServerAPI::SetLatch(uint64_t loopID, uint64_t blockID)
+void PluginServerAPI::SetLatch(LoopOp * loop, mlir::Block* b)
 {
     Json::Value root;
     string funName("SetLatch");
-    root["loopId"] = std::to_string(loopID);
-    root["blockId"] = std::to_string(blockID);
+    root["loopId"] = std::to_string(loop->idAttr().getInt());
+    root["blockId"] = std::to_string(
+        PluginServer::GetInstance()->FindBasicBlock(b));
     string params = root.toStyledString();
     PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
 }
@@ -574,11 +576,12 @@ vector<mlir::Block*> PluginServerAPI::GetLoopBody(uint64_t loopID)
     return BlocksResult(funName, params);
 }
 
-LoopOp PluginServerAPI::GetBlockLoopFather(uint64_t blockID)
+LoopOp PluginServerAPI::GetBlockLoopFather(mlir::Block* b)
 {
     Json::Value root;
     string funName("GetBlockLoopFather");
-    root["blockId"] = std::to_string(blockID);
+    root["blockId"] = std::to_string(
+        PluginServer::GetInstance()->FindBasicBlock(b));
     string params = root.toStyledString();
     return PluginServer::GetInstance()->LoopOpResult(funName, params);
 }
@@ -601,18 +604,17 @@ bool PluginServerAPI::InsertValue(uint64_t id, mlir::Value v)
     return server->InsertValue(id, v);
 }
 
-bool PluginServerAPI::RedirectFallthroughTarget(FallThroughOp& fop,
-                                                uint64_t src, uint64_t dest)
+bool PluginServerAPI::RedirectFallthroughTarget(
+    FallThroughOp& fop, mlir::Block* src, mlir::Block* dest)
 {
     Json::Value root;
     string funName = __func__;
-    root["src"] = src;
-    root["dest"] = dest;
+    root["src"] = PluginServer::GetInstance()->FindBasicBlock(src);
+    root["dest"] = PluginServer::GetInstance()->FindBasicBlock(dest);
     string params = root.toStyledString();
     PluginServer::GetInstance()->RemoteCallClientWithAPI(funName, params);
     // update server
-    PluginServer *server = PluginServer::GetInstance();
-    fop->setSuccessor(server->FindBlock(dest), 0);
+    fop->setSuccessor(dest, 0);
     return true;
 }
 
diff --git a/user/ArrayWidenPass.cpp b/user/ArrayWidenPass.cpp
index 6fb187c..591ecdb 100644
--- a/user/ArrayWidenPass.cpp
+++ b/user/ArrayWidenPass.cpp
@@ -947,8 +947,8 @@ static void update_loop_dominator(uint64_t dir, FunctionOp* funcOp)
         }
         uint64_t i_bbAddr = cfAPI.GetImmediateDominator(dir, bbAddr);
         if (!i_bbAddr || &bb == originLoop.exitBB1) {
-            cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(&bb), \
-                cfAPI.RecomputeDominator(1, pluginAPI.FindBasicBlock(&bb)));
+            cfAPI.SetImmediateDominator(1, &bb, \
+                cfAPI.RecomputeDominator(1, &bb));
             continue;
         }
     }
@@ -962,7 +962,7 @@ static void remove_originLoop(LoopOp *loop, FunctionOp* funcOp)
     body = loop->GetLoopBody();
     unsigned n = loop->numBlockAttr().getInt();
     for (unsigned i = 0; i < n; i++) {
-        controlapi.DeleteBlock(body[i], funcOp->idAttr().getInt(), pluginAPI.FindBasicBlock(body[i]));
+        controlapi.DeleteBlock(body[i], funcOp);
     }
     loop->Delete();
 }
@@ -975,12 +975,11 @@ static void create_prolog_bb(Block *prolog_bb, Block *after_bb, Block *dominator
     ControlFlowAPI cfAPI;
     PluginServerAPI pluginAPI;
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(prolog_bb), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(prolog_bb, outer);
 
     FallThroughOp fp = llvm::dyn_cast<FallThroughOp>(entryEdge.src->back());
-    pluginAPI.RedirectFallthroughTarget(fp, pluginAPI.FindBasicBlock(entryEdge.src),
-        pluginAPI.FindBasicBlock(prolog_bb));
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(prolog_bb), pluginAPI.FindBasicBlock(dominator_bb));
+    pluginAPI.RedirectFallthroughTarget(fp, entryEdge.src, prolog_bb);
+    cfAPI.SetImmediateDominator(1, prolog_bb, dominator_bb);
 
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     lhs1 = baseSsa.Copy();
@@ -993,13 +992,13 @@ static void create_prolog_bb(Block *prolog_bb, Block *after_bb, Block *dominator
         mlir::Attribute attr = opBuilder->getI64IntegerAttr(originLoop.step);
         mlir::Value step = ConstOp::CreateConst(*opBuilder, attr, originLoop.base.getType());
         ops.push_back(step);
-        AssignOp opa = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus, prolog_bb);
+        AssignOp opa = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus);
     } else {
         ops.push_back(lhs1);
         ops.push_back(originLoop.base);
-        AssignOp op = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop, prolog_bb);
+        AssignOp op = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop);
     }
-    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), pluginAPI.FindBasicBlock(prolog_bb), ftBB);
+    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), prolog_bb, ftBB);
     baseSsa.SetCurrentDef(lhs1);
     defs_map.emplace(prolog_bb, lhs1);
 }
@@ -1010,10 +1009,10 @@ static void create_loop_pred_bb(
     ControlFlowAPI cfAPI;
     PluginServerAPI pluginAPI;
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(loop_pred_bb), outer->idAttr().getInt());
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(loop_pred_bb), pluginAPI.FindBasicBlock(dominator_bb));
+    pluginAPI.AddBlockToLoop(loop_pred_bb, outer);
+    cfAPI.SetImmediateDominator(1, loop_pred_bb, dominator_bb);
     opBuilder->setInsertionPointToStart(loop_pred_bb);
-    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), pluginAPI.FindBasicBlock(loop_pred_bb), ftBB);
+    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), loop_pred_bb, ftBB);
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     defs_map.emplace(loop_pred_bb, baseSsa.GetCurrentDef());
 }
@@ -1029,9 +1028,9 @@ static void create_align_loop_header(Block *align_loop_header, Block* after_bb,
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     Value entry_node = baseSsa.GetCurrentDef();
     
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(align_loop_header), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(align_loop_header, outer);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(align_loop_header), pluginAPI.FindBasicBlock(dominator_bb));
+    cfAPI.SetImmediateDominator(1, align_loop_header, dominator_bb);
     phi = PhiOp::CreatePhi(nullptr, align_loop_header);
     cfAPI.CreateNewDef(entry_node, phi.getOperation());
     res = phi.GetResult();
@@ -1043,25 +1042,25 @@ static void create_align_loop_header(Block *align_loop_header, Block* after_bb,
         PluginIR::PluginIntegerType::Unsigned);
     ops.push_back(SSAOp::MakeSSA(*opBuilder, baseType));
     ops.push_back(res);
-    AssignOp op1 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop, (align_loop_header));
+    AssignOp op1 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop);
     Value lhs1 = op1.GetLHS();
     ops.clear();
     ops.push_back(SSAOp::MakeSSA(*opBuilder, lhs1.getType()));
     ops.push_back(lhs1);
     ops.push_back(ConstOp::CreateConst(*opBuilder, opBuilder->getI64IntegerAttr(8), lhs1.getType()));
-    AssignOp op2 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus, (align_loop_header));
+    AssignOp op2 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus);
     Value lhs2 = op2.GetLHS();
     ops.clear();
     ops.push_back(SSAOp::MakeSSA(*opBuilder, baseType));
     ops.push_back(originLoop.limit);
-    AssignOp op3 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop, (align_loop_header));
+    AssignOp op3 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop);
     Value lhs3 = op3.GetLHS();
     ops.clear();
     ops.push_back(lhs2);
     ops.push_back(lhs3);
 
-    cond_stmt = opBuilder->create<CondOp>(opBuilder->getUnknownLoc(), IComparisonCode::le,
-        lhs2, lhs3, tb, fb, (align_loop_header));
+    cond_stmt = opBuilder->create<CondOp>(
+        opBuilder->getUnknownLoc(), IComparisonCode::le, lhs2, lhs3, tb, fb);
 
     baseSsa.SetCurrentDef(res);
     defs_map.emplace(align_loop_header, res);
@@ -1096,9 +1095,9 @@ static LoopOp *init_new_loop(LoopOp *outer_loop, Block* header, Block* latch, Fu
     LoopOp *new_loop;
     LoopOp new_loopt = funcOp->AllocateNewLoop();
     new_loop = &new_loopt;
-    pluginAPI.SetHeader(new_loop->idAttr().getInt(), pluginAPI.FindBasicBlock(header));
-    pluginAPI.SetLatch(new_loop->idAttr().getInt(), pluginAPI.FindBasicBlock(latch));
-    new_loop->AddLoop(outer_loop->idAttr().getInt(), funcOp->idAttr().getInt());
+    pluginAPI.SetHeader(new_loop, header);
+    pluginAPI.SetLatch(new_loop, latch);
+    new_loop->AddLoop(outer_loop, funcOp);
 
     return new_loop;
 }
@@ -1111,10 +1110,9 @@ static void create_align_loop_body_bb(Block *align_loop_body_bb, Block* after_bb
     ControlFlowAPI cfAPI;
     PluginServerAPI pluginAPI;
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(align_loop_body_bb), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(align_loop_body_bb, outer);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(align_loop_body_bb),
-        pluginAPI.FindBasicBlock(dominator_bb));
+    cfAPI.SetImmediateDominator(1, align_loop_body_bb, dominator_bb);
 
     opBuilder->setInsertionPointToStart(align_loop_body_bb);
     llvm::SmallVector<mlir::Value> ops;
@@ -1126,15 +1124,15 @@ static void create_align_loop_body_bb(Block *align_loop_body_bb, Block* after_bb
     ops.push_back(SSAOp::MakeSSA(*opBuilder, sizeType));
 
     ops.push_back(baseSsa.GetCurrentDef());
-    AssignOp op1 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::Nop, (align_loop_body_bb));
+    AssignOp op1 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::Nop);
     Value var = op1.GetLHS();
     ops.clear();
     ops.push_back(SSAOp::MakeSSA(*opBuilder, originLoop.arr2.getType()));
     ops.push_back(originLoop.arr2);
     ops.push_back(var);
-    AssignOp op2 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::PtrPlus, (align_loop_body_bb));
+    AssignOp op2 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::PtrPlus);
     lhs1 = op2.GetLHS();
     ops.clear();
     PluginIR::PluginTypeBase baseType = PluginIR::PluginIntegerType::get(context, 64,
@@ -1143,15 +1141,16 @@ static void create_align_loop_body_bb(Block *align_loop_body_bb, Block* after_bb
     ops.push_back(SSAOp::MakeSSA(*opBuilder, baseType));
     mlir::Attribute asdada = opBuilder->getI64IntegerAttr(0);
     ops.push_back(pluginAPI.BuildMemRef(baseType, lhs1, ConstOp::CreateConst(*opBuilder, asdada, pointerTy)));
-    AssignOp op3 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::UNDEF, (align_loop_body_bb));
+    AssignOp op3 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::UNDEF);
 
     lhs1 = op3.GetLHS();
     ops.clear();
     ops.push_back(SSAOp::MakeSSA(*opBuilder, originLoop.arr1.getType()));
     ops.push_back(originLoop.arr1);
     ops.push_back(var);
-    AssignOp op4 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::PtrPlus, (align_loop_body_bb));
+    AssignOp op4 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::PtrPlus);
     lhs2 = op4.GetLHS();
     ops.clear();
 
@@ -1159,12 +1158,14 @@ static void create_align_loop_body_bb(Block *align_loop_body_bb, Block* after_bb
     ops.push_back(pluginAPI.BuildMemRef(
         baseType, lhs2, ConstOp::CreateConst(*opBuilder, opBuilder->getI64IntegerAttr(0), pointerTy)));
 
-    AssignOp op5 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::UNDEF, (align_loop_body_bb));
+    AssignOp op5 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::UNDEF);
     lhs2 = op5.GetLHS();
 
-    cond_stmt = opBuilder->create<CondOp>(opBuilder->getUnknownLoc(),
-    llvm::dyn_cast<CondOp>(originLoop.condOp2).condCode(), lhs1, lhs2, tb, fb, (align_loop_body_bb));
+    cond_stmt = opBuilder->create<CondOp>(
+        opBuilder->getUnknownLoc(),
+        llvm::dyn_cast<CondOp>(originLoop.condOp2).condCode(),
+        lhs1, lhs2, tb, fb);
 }
 
 static void create_align_loop_latch(Block *align_loop_latch, Block* after_bb, Block* dominator_bb,
@@ -1177,10 +1178,9 @@ static void create_align_loop_latch(Block *align_loop_latch, Block* after_bb, Bl
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     Value entry_node = baseSsa.GetCurrentDef();
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(align_loop_latch), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(align_loop_latch, outer);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(align_loop_latch),
-        pluginAPI.FindBasicBlock(dominator_bb));
+    cfAPI.SetImmediateDominator(1, align_loop_latch, dominator_bb);
 
     res = baseSsa.Copy();
 
@@ -1189,8 +1189,8 @@ static void create_align_loop_latch(Block *align_loop_latch, Block* after_bb, Bl
     ops.push_back(res);
     ops.push_back(entry_node);
     ops.push_back(ConstOp::CreateConst(*opBuilder, opBuilder->getI64IntegerAttr(8), entry_node.getType()));
-    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus, (align_loop_latch));
-    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), pluginAPI.FindBasicBlock(align_loop_latch), ftBB);
+    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus);
+    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), align_loop_latch, ftBB);
     defs_map.emplace(align_loop_latch, res);
 }
 
@@ -1206,10 +1206,9 @@ static void create_align_loop_exit_bb(Block *align_loop_exit_bb, Block* after_bb
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     Value entry_node = baseSsa.GetCurrentDef();
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(align_loop_exit_bb), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(align_loop_exit_bb, outer);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(align_loop_exit_bb),
-        pluginAPI.FindBasicBlock(dominator_bb));
+    cfAPI.SetImmediateDominator(1, align_loop_exit_bb, dominator_bb);
 
     cond_stmt = llvm::dyn_cast<CondOp>(after_bb->back());
     cond_lhs = cond_stmt.GetLHS();
@@ -1220,12 +1219,12 @@ static void create_align_loop_exit_bb(Block *align_loop_exit_bb, Block* after_bb
     ops.push_back((SSAOp::MakeSSA(*opBuilder, cond_lhs.getType())));
     ops.push_back(cond_lhs);
     ops.push_back(cond_rhs);
-    AssignOp op1 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::BitXOR, (align_loop_exit_bb));
+    AssignOp op1 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::BitXOR);
     lhs1 = op1.GetLHS();
     ops.clear();
     ops.push_back(lhs1);
-    build_ctzll = opBuilder->create<CallOp>(opBuilder->getUnknownLoc(), ops, (align_loop_exit_bb));
+    build_ctzll = opBuilder->create<CallOp>(opBuilder->getUnknownLoc(), ops);
     PluginIR::PluginTypeBase intType =
         PluginIR::PluginIntegerType::get(context, 32, PluginIR::PluginIntegerType::Signed);
     lhs1 = SSAOp::MakeSSA(*opBuilder, intType);
@@ -1237,12 +1236,12 @@ static void create_align_loop_exit_bb(Block *align_loop_exit_bb, Block* after_bb
     ops.push_back(lhs2);
     ops.push_back(lhs1);
     ops.push_back(ConstOp::CreateConst(*opBuilder, opBuilder->getI64IntegerAttr(3), lhs1.getType()));
-    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Rshift, (align_loop_exit_bb));
+    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Rshift);
 
     ops.clear();
     ops.push_back(SSAOp::MakeSSA(*opBuilder, entry_node.getType()));
     ops.push_back(lhs2);
-    AssignOp op2 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop, (align_loop_exit_bb));
+    AssignOp op2 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Nop);
     lhs1 = op2.GetLHS();
     SSAOp entrySsa = dyn_cast<mlir::Plugin::SSAOp>(entry_node.getDefiningOp());
     lhs2 = entrySsa.Copy();
@@ -1250,9 +1249,9 @@ static void create_align_loop_exit_bb(Block *align_loop_exit_bb, Block* after_bb
     ops.push_back(lhs2);
     ops.push_back(lhs1);
     ops.push_back(entry_node);
-    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus, (align_loop_exit_bb));
+    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus);
 
-    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), pluginAPI.FindBasicBlock(align_loop_exit_bb), ftBB);
+    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), align_loop_exit_bb, ftBB);
 
     defs_map.emplace(align_loop_exit_bb, lhs2);
 }
@@ -1269,10 +1268,9 @@ static void create_epilogue_loop_header(Block *epilogue_loop_header, Block* afte
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     Value entry_node = baseSsa.GetCurrentDef();
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(epilogue_loop_header), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(epilogue_loop_header, outer);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(epilogue_loop_header),
-        pluginAPI.FindBasicBlock(dominator_bb));
+    cfAPI.SetImmediateDominator(1, epilogue_loop_header, dominator_bb);
 
     phi = PhiOp::CreatePhi(nullptr, epilogue_loop_header);
     cfAPI.CreateNewDef(entry_node, phi.getOperation());
@@ -1280,8 +1278,10 @@ static void create_epilogue_loop_header(Block *epilogue_loop_header, Block* afte
 
     opBuilder->setInsertionPointToStart(epilogue_loop_header);
 
-    cond_stmt = opBuilder->create<CondOp>(opBuilder->getUnknownLoc(),
-        llvm::dyn_cast<CondOp>(originLoop.condOp1).condCode(), res, originLoop.limit, tb, fb, (epilogue_loop_header));
+    cond_stmt = opBuilder->create<CondOp>(
+        opBuilder->getUnknownLoc(),
+        llvm::dyn_cast<CondOp>(originLoop.condOp1).condCode(),
+        res, originLoop.limit, tb, fb);
 
     baseSsa.SetCurrentDef(res);
     defs_map.emplace(epilogue_loop_header, res);
@@ -1299,10 +1299,9 @@ static void create_epilogue_loop_body_bb(Block *epilogue_loop_body_bb, Block* af
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     Value entry_node = baseSsa.GetCurrentDef();
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(epilogue_loop_body_bb), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(epilogue_loop_body_bb, outer);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(epilogue_loop_body_bb),
-        pluginAPI.FindBasicBlock(dominator_bb));
+    cfAPI.SetImmediateDominator(1, epilogue_loop_body_bb, dominator_bb);
 
     opBuilder->setInsertionPointToStart(epilogue_loop_body_bb);
     llvm::SmallVector<mlir::Value> ops;
@@ -1311,15 +1310,15 @@ static void create_epilogue_loop_body_bb(Block *epilogue_loop_body_bb, Block* af
         PluginIR::PluginIntegerType::Unsigned);
     ops.push_back(SSAOp::MakeSSA(*opBuilder, sizeType));
     ops.push_back(entry_node);
-    AssignOp op1 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::Nop, (epilogue_loop_body_bb));
+    AssignOp op1 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::Nop);
     lhs1 = op1.GetLHS();
     ops.clear();
     ops.push_back(SSAOp::MakeSSA(*opBuilder, originLoop.arr1.getType()));
     ops.push_back(originLoop.arr1);
     ops.push_back(lhs1);
-    AssignOp op2 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::PtrPlus, (epilogue_loop_body_bb));
+    AssignOp op2 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::PtrPlus);
     lhs2 = op2.GetLHS();
 
     ops.clear();
@@ -1328,7 +1327,7 @@ static void create_epilogue_loop_body_bb(Block *epilogue_loop_body_bb, Block* af
     ops.push_back(SSAOp::MakeSSA(*opBuilder, charType));
     ops.push_back(pluginAPI.BuildMemRef(charType, lhs2,
         ConstOp::CreateConst(*opBuilder, opBuilder->getI64IntegerAttr(0), lhs2.getType())));
-    g = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::UNDEF, (epilogue_loop_body_bb));
+    g = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::UNDEF);
 
     lhs2 = g.GetLHS();
 
@@ -1336,8 +1335,8 @@ static void create_epilogue_loop_body_bb(Block *epilogue_loop_body_bb, Block* af
     ops.push_back(SSAOp::MakeSSA(*opBuilder, originLoop.arr2.getType()));
     ops.push_back(originLoop.arr2);
     ops.push_back(lhs1);
-    AssignOp op3 = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops,
-        IExprCode::PtrPlus, (epilogue_loop_body_bb));
+    AssignOp op3 = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::PtrPlus);
     lhs3 = op3.GetLHS();
 
     ops.clear();
@@ -1345,11 +1344,14 @@ static void create_epilogue_loop_body_bb(Block *epilogue_loop_body_bb, Block* af
     ops.push_back(SSAOp::MakeSSA(*opBuilder, charType));
     ops.push_back(pluginAPI.BuildMemRef(charType, lhs3, ConstOp::CreateConst(*opBuilder,
         opBuilder->getI64IntegerAttr(0), lhs3.getType())));
-    g = opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::UNDEF, (epilogue_loop_body_bb));
+    g = opBuilder->create<AssignOp>(
+        opBuilder->getUnknownLoc(), ops, IExprCode::UNDEF);
 
     Value res = g.GetLHS();
-    cond_stmt = opBuilder->create<CondOp>(opBuilder->getUnknownLoc(),
-        llvm::dyn_cast<CondOp>(originLoop.condOp1).condCode(), lhs2, res, tb, fb, (epilogue_loop_body_bb));
+    cond_stmt = opBuilder->create<CondOp>(
+        opBuilder->getUnknownLoc(),
+        llvm::dyn_cast<CondOp>(originLoop.condOp1).condCode(),
+        lhs2, res, tb, fb);
 
     defs_map.emplace(epilogue_loop_body_bb, baseSsa.GetCurrentDef());
 }
@@ -1363,10 +1365,9 @@ static void create_epilogue_loop_latch(Block *epilogue_loop_latch,
     SSAOp baseSsa = dyn_cast<mlir::Plugin::SSAOp>(originLoop.base.getDefiningOp());
     Value entry_node = baseSsa.GetCurrentDef();
 
-    pluginAPI.AddBlockToLoop(pluginAPI.FindBasicBlock(epilogue_loop_latch), outer->idAttr().getInt());
+    pluginAPI.AddBlockToLoop(epilogue_loop_latch, outer);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(epilogue_loop_latch),
-        pluginAPI.FindBasicBlock(dominator_bb));
+    cfAPI.SetImmediateDominator(1, epilogue_loop_latch, dominator_bb);
 
     SSAOp entrySsa = dyn_cast<mlir::Plugin::SSAOp>(entry_node.getDefiningOp());
     res = entrySsa.Copy();
@@ -1377,9 +1378,10 @@ static void create_epilogue_loop_latch(Block *epilogue_loop_latch,
     ops.push_back(entry_node);
     ops.push_back(ConstOp::CreateConst(*opBuilder,
         opBuilder->getI64IntegerAttr(originLoop.step), entry_node.getType()));
-    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), ops, IExprCode::Plus, (epilogue_loop_latch));
+    opBuilder->create<AssignOp>(opBuilder->getUnknownLoc(), 
+                                ops, IExprCode::Plus);
 
-    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), pluginAPI.FindBasicBlock(epilogue_loop_latch), ftBB);
+    opBuilder->create<FallThroughOp>(opBuilder->getUnknownLoc(), epilogue_loop_latch, ftBB);
 
     defs_map.emplace(epilogue_loop_latch, res);
 }
@@ -1395,47 +1397,27 @@ static void create_new_loops(edge entryEdge, FunctionOp *funcOp)
     LoopOp *epilogue_loop;
     PluginServerAPI pluginAPI;
     ControlFlowAPI cfAPI;
-    LoopOp outer = pluginAPI.GetBlockLoopFather(pluginAPI.FindBasicBlock(entryEdge.src));
+    LoopOp outer = pluginAPI.GetBlockLoopFather(entryEdge.src);
 
-    uint64_t bbAddr = cfAPI.CreateBlock(entryEdge.src, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(entryEdge.src));
-    prolog_bb = pluginAPI.FindBlock(bbAddr);
+    prolog_bb = cfAPI.CreateBlock(entryEdge.src, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(prolog_bb, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(prolog_bb));
-    align_pred_bb = pluginAPI.FindBlock(bbAddr);
+    align_pred_bb = cfAPI.CreateBlock(prolog_bb, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(align_pred_bb, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(align_pred_bb));
-    align_loop_header = pluginAPI.FindBlock(bbAddr);
+    align_loop_header = cfAPI.CreateBlock(align_pred_bb, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(align_loop_header, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(align_loop_header));
-    align_loop_body_bb = pluginAPI.FindBlock(bbAddr);
+    align_loop_body_bb = cfAPI.CreateBlock(align_loop_header, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(align_loop_body_bb, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(align_loop_body_bb));
-    align_loop_latch = pluginAPI.FindBlock(bbAddr);
+    align_loop_latch = cfAPI.CreateBlock(align_loop_body_bb, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(align_loop_body_bb, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(align_loop_body_bb));
-    align_loop_exit_bb = pluginAPI.FindBlock(bbAddr);
+    align_loop_exit_bb = cfAPI.CreateBlock(align_loop_body_bb, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(align_loop_header, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(align_loop_header));
-    epilogue_loop_pred_bb = pluginAPI.FindBlock(bbAddr);
+    epilogue_loop_pred_bb = cfAPI.CreateBlock(align_loop_header, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(epilogue_loop_pred_bb, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(epilogue_loop_pred_bb));
-    epilogue_loop_header = pluginAPI.FindBlock(bbAddr);
+    epilogue_loop_header = cfAPI.CreateBlock(epilogue_loop_pred_bb, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(epilogue_loop_header, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(epilogue_loop_header));
-    epilogue_loop_body_bb = pluginAPI.FindBlock(bbAddr);
+    epilogue_loop_body_bb = cfAPI.CreateBlock(epilogue_loop_header, funcOp);
 
-    bbAddr = cfAPI.CreateBlock(epilogue_loop_body_bb, funcOp->idAttr().getInt(),
-        pluginAPI.FindBasicBlock(epilogue_loop_body_bb));
-    epilogue_loop_latch = pluginAPI.FindBlock(bbAddr);
+    epilogue_loop_latch = cfAPI.CreateBlock(epilogue_loop_body_bb, funcOp);
 
     create_prolog_bb(prolog_bb, entryEdge.src, entryEdge.src, &outer, entryEdge, funcOp, align_pred_bb);
     
@@ -1473,18 +1455,14 @@ static void create_new_loops(edge entryEdge, FunctionOp *funcOp)
 
     epilogue_loop = init_new_loop(&outer, epilogue_loop_header, epilogue_loop_latch, funcOp);
 
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(originLoop.exitBB1),
-        pluginAPI.FindBasicBlock(entryEdge.src));
-    cfAPI.SetImmediateDominator(1, pluginAPI.FindBasicBlock(originLoop.exitBB2),
-        pluginAPI.FindBasicBlock(entryEdge.src));
+    cfAPI.SetImmediateDominator(1, originLoop.exitBB1, entryEdge.src);
+    cfAPI.SetImmediateDominator(1, originLoop.exitBB2, entryEdge.src);
 
     rewrite_add_phi_arg(originLoop.exitBB1);
     rewrite_add_phi_arg(originLoop.exitBB2);
 
-    cfAPI.RemoveEdge(pluginAPI.FindBasicBlock(originLoop.exitE1.src),
-        pluginAPI.FindBasicBlock(originLoop.exitE1.dest));
-    cfAPI.RemoveEdge(pluginAPI.FindBasicBlock(originLoop.exitE2.src),
-        pluginAPI.FindBasicBlock(originLoop.exitE2.dest));
+    cfAPI.RemoveEdge(originLoop.exitE1.src, originLoop.exitE1.dest);
+    cfAPI.RemoveEdge(originLoop.exitE2.src, originLoop.exitE2.dest);
 }
 
 static void convertToNewLoop(LoopOp* loop, FunctionOp* funcOp)
diff --git a/user/LocalVarSummeryPass.cpp b/user/LocalVarSummeryPass.cpp
index ccee9f7..a343175 100755
--- a/user/LocalVarSummeryPass.cpp
+++ b/user/LocalVarSummeryPass.cpp
@@ -69,7 +69,7 @@ static void LocalVarSummery(void)
             size_t paramIndex = 0;
             llvm::ArrayRef<mlir::Type> paramsType = ty.getParams();
             for (auto ty : ty.getParams()) {
-                printf("\n    Param index : %d\n", paramIndex++);
+                printf("\n    Param index : %ld\n", paramIndex++);
                 printf("\n    Param type id : %d\n", ty.dyn_cast<PluginIR::PluginTypeBase>().getPluginTypeID());
             }
         }
-- 
2.33.0