[LowerToAIE] Fix for non-deterministic error (observed with convoluti…

…on) (#694)

What was happening was 

1) Value/Operation inserted into IRMapping
2) IRRewriter erased Value/Op
3) New Value/Op created with same address as erased Value/Op
4) Bad state: new Value/Op is a key in IRMapping, even though it was
never added.

That's my analysis, at least. Either way, this change makes my failure
rate go from 10/20 to 0/20.

An alternative to this PR might be to postpone erasing of ops even
longer, so that only after all CoreOps have been processed are the ops
erased.

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELowerToAIE.cpp

@@ -19,7 +19,6 @@
 #include "iree-amd-aie/IR/AMDAIEOps.h"
 #include "iree-amd-aie/Transforms/AMDAIEUtils.h"
 #include "iree-amd-aie/Transforms/Passes.h"
-#include "iree-amd-aie/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Iterators.h"
@@ -43,12 +42,26 @@ void remapOperands(Operation *op, IRMapping &mapper) {
   }
 }
 
+/// It is dangerous to erase ops with `rewriter` without erasing them from
+/// `mapper` too, as addresses of Operations/Values can be reused, resulting in
+/// unexpected key-value pairs in `mapper`. Use this utility if `mapper` might
+/// be used after `op` is erased.
+void eraseOp(IRRewriter &rewriter, IRMapping &mapper, Operation *op) {
+  for (Value result : op->getResults()) {
+    mapper.erase(result);
+  }
+  mapper.erase(op);
+  op->dropAllUses();
+  rewriter.eraseOp(op);
+}
+
 //===----------------------------------------------------------------------===//
 // Convert amdaie.core operation to aie.core
 //===----------------------------------------------------------------------===//
 
 namespace {
 
+
 /// Utility to convert vectors of `size` and `stride` into an
 /// `AIE::BDDimLayoutArrayAttr`.
 AIE::BDDimLayoutArrayAttr convertSizeStrideToBDDimLayoutArrayAttr(
@@ -259,7 +272,7 @@ LogicalResult coreLinalgOpToAIE(IRRewriter &rewriter, linalg::LinalgOp linalgOp,
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(linalgOp);
   rewriter.clone(*(linalgOp.getOperation()), mapper);
-  rewriter.eraseOp(linalgOp);
+  eraseOp(rewriter, mapper, linalgOp);
   return success();
 }
 
@@ -422,8 +435,7 @@ LogicalResult coreToAIE(IRRewriter &rewriter, AMDAIE::CoreOp coreOp,
     return failure();
   }
   for (auto *op : toBeErased) {
-    op->dropAllUses();
-    rewriter.eraseOp(op);
+    eraseOp(rewriter, mapper, op);
   }
 
   mapper.map(coreOp.getResult(), aieCoreOp.getResult());
@@ -687,7 +699,7 @@ LogicalResult controlCodeToAie(IRRewriter &rewriter,
                                          bindingsMapper);
                 })
                 .Case<AMDAIE::EndOp>([&](auto endOp) {
-                  rewriter.eraseOp(endOp);
+                  eraseOp(rewriter, mapper, endOp);
                   return success();
                 })
                 .Default([&](Operation *op) {
@@ -701,8 +713,7 @@ LogicalResult controlCodeToAie(IRRewriter &rewriter,
       });
   if (res.wasInterrupted()) return failure();
   for (auto *op : toBeErased) {
-    op->dropAllUses();
-    rewriter.eraseOp(op);
+    eraseOp(rewriter, mapper, op);
   }
   return success();
 }
@@ -917,7 +928,8 @@ LogicalResult lowerToAIE(ModuleOp moduleOp) {
     rewriter.moveOpBefore(ipuFuncOp, deviceBlock, deviceBlock->end());
     // After walking the FuncOp, it has been converted into a DeviceOp and can
     // safely be erased.
-    rewriter.eraseOp(funcOp);
+    eraseOp(rewriter, mapper, funcOp);
+
     return WalkResult::advance();
   });
   if (funcRes.wasInterrupted()) return failure();