Fix the deallocation problem for sequence op

src/Compiler/CompilerPasses.cpp

@@ -156,9 +156,9 @@ void addKrnlToLLVMPasses(
   }
 
   pm.addNestedPass<func::FuncOp>(krnl::createLowerKrnlRegionPass());
+  pm.addNestedPass<func::FuncOp>(krnl::createConvertSeqToMemrefPass());
   pm.addNestedPass<func::FuncOp>(mlir::createConvertSCFToCFPass());
 
-  pm.addNestedPass<func::FuncOp>(krnl::createConvertSeqToMemrefPass());
   pm.addPass(krnl::createConvertKrnlToLLVMPass(verifyInputTensors));
   pm.addPass(mlir::createReconcileUnrealizedCastsPass());
   pm.addPass(mlir::createCanonicalizerPass());

src/Conversion/ONNXToKrnl/ControlFlow/Loop.cpp

@@ -311,8 +311,9 @@ struct ONNXLoopOpLowering : public ConversionPattern {
                 rewriter.setInsertionPointToStart(
                     &regionOp.bodyRegion().front());
                 Value origIV = loopInd[0];
-                auto src = rewriter.create<KrnlSeqExtractOp>(
-                    loc, seqElementType, output, origIV);
+                auto src = rewriter.create<KrnlSeqExtractOp>(loc,
+                    seqElementType, output, origIV,
+                    IntegerAttr::get(rewriter.getIntegerType(1, false), 0));
                 emitCopy(rewriter, loc, src, alloc, {origIV});
               });
           newOutputs.emplace_back(alloc);

src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp

@@ -347,8 +347,9 @@ void FrontendToKrnlLoweringPass::runOnOperation() {
 
   // If `emitDealloc` is turned off, make sure we don't have buffer deallocation
   // at this level. Will use MLIR buffer-deallocation for this purpose instead.
-  if (!emitDealloc)
-    target.addIllegalOp<mlir::memref::DeallocOp>();
+  // However, since the SequenceErase needs to emit memref dealloc, the previous
+  // the following statement is commented out (Chentong)
+  // if (!emitDealloc) target.addIllegalOp<mlir::memref::DeallocOp>();
 
   // TODO: enable this once more ops are supported.
   // We also define the ONNX dialect as Illegal so that the conversion will fail

src/Conversion/ONNXToKrnl/Sequence/SequenceAt.cpp

@@ -31,17 +31,21 @@ struct ONNXSequenceAtOpLowering : public ConversionPattern {
     MultiDialectBuilder<KrnlBuilder, MemRefBuilder> create(rewriter, loc);
     IndexExprScope IEScope(&rewriter, loc);
 
-    auto input_sequence = operandAdaptor.input_sequence();
+    Value input_sequence = operandAdaptor.input_sequence();
+    Type outputMemRefType =
+        input_sequence.getType().cast<MemRefType>().getElementType();
     auto dimSize = create.mem.dim(input_sequence, 0);
     SymbolIndexExpr boundIE(dimSize);
     IndexExpr positionIE =
         SymbolIndexExpr(create.krnl.load(operandAdaptor.position()));
-
-    // Just for lit test to pass
-    auto correctionIE = positionIE + boundIE;
-    positionIE = IndexExpr::select(positionIE < 0, correctionIE, positionIE);
-    auto outputVal = create.krnl.load(
-        operandAdaptor.input_sequence(), positionIE.getValue());
+    // Handle the negative position
+    IndexExpr condIE = positionIE < 0;
+    IndexExpr fixedPosition = positionIE + boundIE;
+    positionIE = IndexExpr::select(condIE, fixedPosition, positionIE);
+
+    Value outputVal = rewriter.create<KrnlSeqExtractOp>(loc, outputMemRefType,
+        input_sequence, positionIE.getValue(),
+        IntegerAttr::get(rewriter.getIntegerType(1, false), 1));
 
     rewriter.replaceOp(op, outputVal);
     return success();

src/Conversion/ONNXToKrnl/Sequence/SequenceErase.cpp

@@ -28,7 +28,7 @@ struct ONNXSequenceEraseOpLowering : public ConversionPattern {
       ConversionPatternRewriter &rewriter) const final {
     Location loc = op->getLoc();
     ONNXSequenceEraseOpAdaptor operandAdaptor(operands);
-    ONNXSequenceInsertOp thisOp = dyn_cast<ONNXSequenceInsertOp>(op);
+    ONNXSequenceEraseOp thisOp = dyn_cast<ONNXSequenceEraseOp>(op);
     MultiDialectBuilder<MathBuilder, MemRefBuilder> create(rewriter, loc);
     IndexExprScope IEScope(&rewriter, loc);
 
@@ -66,7 +66,7 @@ struct ONNXSequenceEraseOpLowering : public ConversionPattern {
       positionIE = IndexExpr::select(positionIE < 0, correctionIE, positionIE);
     }
 
-    // Copy before the insert
+    // Copy the elements before the position
     KrnlBuilder createKrnl(rewriter, loc);
     SmallVector<IndexExpr, 1> lbs;
     lbs.emplace_back(LiteralIndexExpr(0));
@@ -80,9 +80,12 @@ struct ONNXSequenceEraseOpLowering : public ConversionPattern {
           createKrnl.store(element, alloc, indicesLoopInd[0]);
         });
 
-    // ToDo (chentong)Free the erased element
+    // Free the element to be erased
+    Value element =
+        createKrnl.load(operandAdaptor.input_sequence(), positionIE.getValue());
+    create.mem.dealloc(element);
 
-    // Copy after the insert
+    // Copy the elements after the position
     SmallVector<IndexExpr, 1> lbs1;
     lbs1.emplace_back(positionIE + 1);
     SmallVector<IndexExpr, 1> ubs1;

src/Conversion/ONNXToKrnl/Sequence/SequenceInsert.cpp

@@ -40,21 +40,11 @@ struct ONNXSequenceInsertOpLowering : public ConversionPattern {
            "Failed to convert type to MemRefType");
     MemRefType outputMemRefType = convertedType.cast<MemRefType>();
 
-    auto seqElementConvertedType =
-        outputMemRefType.getElementType().cast<MemRefType>();
     auto input_sequence = operandAdaptor.input_sequence();
     auto dimSize = create.mem.dim(input_sequence, 0);
     SymbolIndexExpr boundIE(dimSize);
 
-    // Output sequence has one more element
-    auto outputBound = boundIE + 1;
-    SmallVector<IndexExpr, 1> ubsIE;
-    ubsIE.emplace_back(outputBound);
-    Value alloc =
-        insertAllocAndDeallocSimple(rewriter, op, outputMemRefType, loc, ubsIE);
-
-    // Fill the output sequence
-
+    // Handle Optional and negative position
     IndexExpr positionIE;
     if (isFromNone(operandAdaptor.position())) {
       // Insert at the end of the sequence
@@ -65,45 +55,14 @@ struct ONNXSequenceInsertOpLowering : public ConversionPattern {
     } else {
       positionIE = SymbolIndexExpr(create.krnl.load(operandAdaptor.position()));
       // Handle the negative position
-      positionIE =
-          IndexExpr::select(positionIE < 0, positionIE + boundIE, positionIE);
+      IndexExpr condIE = positionIE < 0;
+      IndexExpr fixedPosition = positionIE + boundIE;
+      positionIE = IndexExpr::select(condIE, fixedPosition, positionIE);
     }
 
-    // Copy elements before the insertion position
-    SmallVector<IndexExpr, 1> lbs;
-    lbs.emplace_back(LiteralIndexExpr(0));
-    SmallVector<IndexExpr, 1> ubs;
-    ubs.emplace_back(positionIE);
-    ValueRange firstLoopDef = create.krnl.defineLoops(1);
-    create.krnl.iterateIE(firstLoopDef, firstLoopDef, lbs, ubs,
-        [&](KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
-          auto element = createKrnl.load(
-              operandAdaptor.input_sequence(), indicesLoopInd[0]);
-          auto converted = create.mem.cast(element, seqElementConvertedType);
-          createKrnl.store(converted, alloc, indicesLoopInd[0]);
-        });
-
-    // Insert the input tensor
-    // ToDo (chentong): need to duplicate the tensor
-    auto element =
-        create.mem.cast(operandAdaptor.tensor(), seqElementConvertedType);
-    create.krnl.store(element, alloc, positionIE.getValue());
-
-    // Copy elements after the insertion position
-    SmallVector<IndexExpr, 1> lbs1;
-    lbs1.emplace_back(positionIE + 1);
-    SmallVector<IndexExpr, 1> ubs1;
-    ubs1.emplace_back(outputBound);
-    ValueRange secondLoopDef = create.krnl.defineLoops(1);
-    create.krnl.iterateIE(secondLoopDef, secondLoopDef, lbs1, ubs1,
-        [&](KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
-          auto element = createKrnl.load(
-              operandAdaptor.input_sequence(), indicesLoopInd[0]);
-          auto converted = create.mem.cast(element, seqElementConvertedType);
-          auto outputIndex =
-              create.math.add(indicesLoopInd[0], create.math.constantIndex(1));
-          createKrnl.store(converted, alloc, outputIndex);
-        });
+    Value alloc = rewriter.create<KrnlSeqInsertOp>(loc, outputMemRefType,
+        operandAdaptor.tensor(), operandAdaptor.input_sequence(),
+        positionIE.getValue());
 
     rewriter.replaceOp(op, alloc);
     return success();

src/Conversion/SeqToMemref/CMakeLists.txt

@@ -3,6 +3,7 @@
 add_onnx_mlir_library(OMSeqToMemref
   ConvertSeqToMemref.cpp
   KrnlSeqExtract.cpp
+  KrnlSeqInsert.cpp
   KrnlSeqStore.cpp
 
   LINK_LIBS PUBLIC

src/Conversion/SeqToMemref/ConvertSeqToMemref.cpp

@@ -59,10 +59,11 @@ void ConvertSeqToMemrefPass::runOnOperation() {
   ConversionTarget target(getContext());
 
   target.addIllegalOp<KrnlSeqExtractOp>();
+  target.addIllegalOp<KrnlSeqInsertOp>();
   target.addIllegalOp<KrnlSeqStoreOp>();
   target.addLegalDialect<mlir::AffineDialect, mlir::arith::ArithmeticDialect,
       mlir::memref::MemRefDialect, mlir::func::FuncDialect,
-      mlir::vector::VectorDialect>();
+      mlir::vector::VectorDialect, mlir::scf::SCFDialect>();
 
   // Now that the conversion target has been defined, we just need to provide
   // the set of patterns that will lower the frontend operations.
@@ -71,6 +72,7 @@ void ConvertSeqToMemrefPass::runOnOperation() {
   // Define patterns.
   KrnlTypeConverter typeConverter;
   populateLoweringKrnlSeqExtractOpPattern(typeConverter, patterns, ctx);
+  populateLoweringKrnlSeqInsertOpPattern(typeConverter, patterns, ctx);
   populateLoweringKrnlSeqStoreOpPattern(typeConverter, patterns, ctx);
 
   if (failed(applyPartialConversion(

src/Conversion/SeqToMemref/ConvertSeqToMemref.hpp

@@ -24,6 +24,9 @@ namespace krnl {
 void populateLoweringKrnlSeqExtractOpPattern(mlir::TypeConverter &typeConverter,
     mlir::RewritePatternSet &patterns, mlir::MLIRContext *ctx);
 
+void populateLoweringKrnlSeqInsertOpPattern(mlir::TypeConverter &typeConverter,
+    mlir::RewritePatternSet &patterns, mlir::MLIRContext *ctx);
+
 void populateLoweringKrnlSeqStoreOpPattern(mlir::TypeConverter &typeConverter,
     mlir::RewritePatternSet &patterns, mlir::MLIRContext *ctx);
 

src/Conversion/SeqToMemref/KrnlSeqExtract.cpp

@@ -20,6 +20,7 @@
 #include "src/Conversion/KrnlToLLVM/KrnlToLLVMHelper.hpp"
 #include "src/Dialect/Krnl/KrnlHelper.hpp"
 #include "src/Dialect/Krnl/KrnlOps.hpp"
+#include "src/Dialect/Mlir/DialectBuilder.hpp"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "krnl_to_llvm"
@@ -40,7 +41,9 @@ class KrnlSeqExtractOpLowering : public ConversionPattern {
   LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands,
       ConversionPatternRewriter &rewriter) const override {
     KrnlSeqExtractOpAdaptor operandAdaptor(operands);
+    KrnlSeqExtractOp thisOp = dyn_cast<KrnlSeqExtractOp>(op);
     auto loc = op->getLoc();
+    MultiDialectBuilder<MathBuilder, MemRefBuilder> create(rewriter, loc);
 
     auto output = rewriter
                       .create<memref::LoadOp>(
@@ -51,8 +54,25 @@ class KrnlSeqExtractOpLowering : public ConversionPattern {
     // if the element is read from seq after extracted, or deep deallocation
     // is added when seq is freed
 
-    rewriter.replaceOp(op, output);
-    return success();
+    if (thisOp.copy() == 0) {
+      rewriter.replaceOp(op, output);
+      return success();
+    } else {
+      if (!output.getType().isa<MemRefType>())
+        llvm_unreachable(
+            "Not implemented: type of onnx seq element is not tensor");
+      auto outputType = output.getType().cast<MemRefType>();
+      SmallVector<mlir::Value, 4> allocParams;
+      for (size_t i = 0; i < outputType.getShape().size(); i++) {
+        if (outputType.getShape()[i] == -1) {
+          allocParams.emplace_back(create.mem.dim(output, i));
+        }
+      }
+      Value alloc = create.mem.alignedAlloc(outputType, allocParams);
+      rewriter.create<memref::CopyOp>(loc, output, alloc);
+      rewriter.replaceOp(op, alloc);
+      return success();
+    }
   }
 };