Proper increfs/decrefs for input arrays

Hardcode84 · Mar 3, 2021 · Feb 26, 2021 · Feb 26, 2021 · Feb 26, 2021 · Feb 26, 2021
commit f743c0de85a37ef52ec58697e2e6e98c6e02a62c
diff --git a/mlir-compiler/mlir-compiler/src/pipelines/lower_to_llvm.cpp b/mlir-compiler/mlir-compiler/src/pipelines/lower_to_llvm.cpp
@@ -610,12 +610,6 @@ struct ApplyFastmathFlags : public mlir::OpRewritePattern<Op>
 };
 
 // Copypaste from StandardToLLVM
-mlir::Value createIndexAttrConstant(mlir::OpBuilder &builder, mlir::Location loc,
-                                     mlir::Type resultType, int64_t value) {
-    return builder.create<mlir::LLVM::ConstantOp>(
-        loc, resultType, builder.getIntegerAttr(builder.getIndexType(), value));
-}
-
 struct AllocLikeOpLowering : public mlir::ConvertToLLVMPattern {
     using ConvertToLLVMPattern::createIndexConstant;
     using ConvertToLLVMPattern::getIndexType;
@@ -625,19 +619,6 @@ struct AllocLikeOpLowering : public mlir::ConvertToLLVMPattern {
         : ConvertToLLVMPattern(opName, &converter.getContext(), converter, /*benefit*/99) {}
 
 protected:
-    // Returns 'input' aligned up to 'alignment'. Computes
-    // bumped = input + alignement - 1
-    // aligned = bumped - bumped % alignment
-//    static mlir::Value createAligned(mlir::ConversionPatternRewriter &rewriter, mlir::Location loc,
-//                               mlir::Value input, mlir::Value alignment) {
-//        using namespace mlir;
-//        Value one = createIndexAttrConstant(rewriter, loc, alignment.getType(), 1);
-//        Value bump = rewriter.create<LLVM::SubOp>(loc, alignment, one);
-//        Value bumped = rewriter.create<LLVM::AddOp>(loc, input, bump);
-//        Value mod = rewriter.create<LLVM::URemOp>(loc, bumped, alignment);
-//        return rewriter.create<LLVM::SubOp>(loc, bumped, mod);
-//    }
-
     // Creates a call to an allocation function with params and casts the
     // resulting void pointer to ptrType.
     mlir::Value createAllocCall(mlir::Location loc, mlir::StringRef name, mlir::Type ptrType,
@@ -1227,6 +1208,51 @@ struct PostLLVMLowering :
     }
 };
 
+struct LowerRetain : public mlir::OpConversionPattern<plier::RetainOp>
+{
+    using mlir::OpConversionPattern<plier::RetainOp>::OpConversionPattern;
+
+    mlir::LogicalResult
+    matchAndRewrite(plier::RetainOp op, llvm::ArrayRef<mlir::Value> operands,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+        assert(operands.size() == 1);
+        auto arg = operands[0];
+        if (!arg.getType().isa<mlir::LLVM::LLVMStructType>())
+        {
+            return mlir::failure();
+        }
+
+        auto llvmVoidPointerType =
+            mlir::LLVM::LLVMPointerType::get(rewriter.getIntegerType(8));
+        auto incref_func = [&]()
+        {
+            auto mod = op->getParentOfType<mlir::ModuleOp>();
+            assert(mod);
+            auto func = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>("NRT_incref");
+            if (!func)
+            {
+                mlir::OpBuilder::InsertionGuard guard(rewriter);
+                rewriter.setInsertionPointToStart(mod.getBody());
+                auto llvmVoidType = mlir::LLVM::LLVMVoidType::get(rewriter.getContext());
+                func = rewriter.create<mlir::LLVM::LLVMFuncOp>(
+                    rewriter.getUnknownLoc(), "NRT_incref",
+                    mlir::LLVM::LLVMFunctionType::get(llvmVoidType, llvmVoidPointerType));
+            }
+            return func;
+        }();
+
+        auto loc = op.getLoc();
+        auto index = rewriter.getI64ArrayAttr(0);
+        auto elemType = arg.getType().cast<mlir::LLVM::LLVMStructType>().getBody()[0];
+        mlir::Value ptr = rewriter.create<mlir::LLVM::ExtractValueOp>(loc, elemType, arg, index);
+        ptr = rewriter.create<mlir::LLVM::BitcastOp>(loc, llvmVoidPointerType, ptr);
+        rewriter.create<mlir::LLVM::CallOp>(loc, incref_func, ptr);
+        rewriter.replaceOp(op, arg);
+
+        return mlir::success();
+    }
+};
+
 struct LowerCasts : public mlir::OpConversionPattern<plier::CastOp>
 {
     using mlir::OpConversionPattern<plier::CastOp>::OpConversionPattern;
@@ -1277,7 +1303,7 @@ struct LLVMLoweringPass : public mlir::PassWrapper<LLVMLoweringPass, mlir::Opera
 
     OwningRewritePatternList patterns;
     populateStdToLLVMConversionPatterns(typeConverter, patterns);
-    patterns.insert<LowerCasts>(typeConverter, &getContext());
+    patterns.insert<LowerCasts, LowerRetain>(typeConverter, &getContext());
     patterns.insert<AllocOpLowering, DeallocOpLowering>(typeConverter);
 
     LLVMConversionTarget target(getContext());

diff --git a/mlir-compiler/mlir-compiler/src/pipelines/plier_to_linalg.cpp b/mlir-compiler/mlir-compiler/src/pipelines/plier_to_linalg.cpp
@@ -859,6 +859,41 @@ struct LoopInvariantCodeMotion : public mlir::OpRewritePattern<mlir::scf::ForOp>
     }
 };
 
+struct RetainArgsPass :
+    public mlir::PassWrapper<RetainArgsPass, mlir::FunctionPass>
+{
+    virtual void getDependentDialects(
+        mlir::DialectRegistry &registry) const override
+    {
+        registry.insert<plier::PlierDialect>();
+    }
+
+    void runOnFunction() override;
+};
+
+void RetainArgsPass::runOnFunction()
+{
+    auto func = getFunction();
+    if (func.isPrivate() || func.isDeclaration() || func.body().empty())
+    {
+        return;
+    }
+
+    mlir::OpBuilder builder(&getContext());
+    auto loc = builder.getUnknownLoc();
+    auto block = &func.body().front();
+    builder.setInsertionPointToStart(block);
+    for (auto arg : block->getArguments())
+    {
+        if (arg.getType().isa<mlir::MemRefType>())
+        {
+            auto retained = builder.create<plier::RetainOp>(loc, arg);
+            llvm::SmallPtrSet<mlir::Operation*, 1> except({retained});
+            arg.replaceAllUsesExcept(retained, except);
+        }
+    }
+}
+
 struct PostLinalgOptPass :
     public mlir::PassWrapper<PostLinalgOptPass, mlir::OperationPass<mlir::ModuleOp>>
 {
@@ -921,6 +956,7 @@ void populate_plier_to_linalg_opt_pipeline(mlir::OpPassManager& pm)
     pm.addNestedPass<mlir::FuncOp>(mlir::createBufferLoopHoistingPass());
     pm.addNestedPass<mlir::FuncOp>(mlir::createPromoteBuffersToStackPass());
 
+    pm.addNestedPass<mlir::FuncOp>(std::make_unique<RetainArgsPass>());
     pm.addNestedPass<mlir::FuncOp>(mlir::createBufferDeallocationPass());
 
     pm.addPass(std::make_unique<LowerLinalgPass>());

diff --git a/mlir-compiler/plier/include/plier/PlierOps.td b/mlir-compiler/plier/include/plier/PlierOps.td
@@ -214,6 +214,16 @@ def GetattrOp : Plier_Op<"getattr", [NoSideEffect]> {
     ];
 }
 
+def RetainOp : Plier_Op<"retain"> {
+    let arguments = (ins AnyMemRef:$value);
+
+    let results = (outs Res<AnyMemRef, "", [MemAlloc<DefaultResource>]>:$memref);
+
+    let builders = [
+        OpBuilderDAG<(ins "::mlir::Value":$value)>
+    ];
+}
+
 def ParallelOp : Plier_Op<"parallel",
       [AttrSizedOperandSegments,
        DeclareOpInterfaceMethods<LoopLikeOpInterface>,

diff --git a/mlir-compiler/plier/include/plier/dialect.hpp b/mlir-compiler/plier/include/plier/dialect.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <mlir/IR/BuiltinTypes.h>
 #include <mlir/IR/Dialect.h>
 #include <mlir/IR/Types.h>
 #include <mlir/IR/OpDefinition.h>
@@ -14,6 +15,7 @@ using Value = ::mlir::Value;
 using Region = ::mlir::Region;
 using LogicalResult = ::mlir::LogicalResult;
 using Operation = ::mlir::Operation;
+namespace MemoryEffects = ::mlir::MemoryEffects;
 
 template<typename T>
 using ArrayRef = ::mlir::ArrayRef<T>;

diff --git a/mlir-compiler/plier/src/dialect.cpp b/mlir-compiler/plier/src/dialect.cpp
@@ -336,6 +336,11 @@ void GetattrOp::getCanonicalizationPatterns(
     results.insert<GetattrGlobalRewrite>(context);
 }
 
+void RetainOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
+                     mlir::Value value) {
+    RetainOp::build(builder, state, value.getType(), value);
+}
+
 mlir::LogicalResult ParallelOp::moveOutOfLoop(mlir::ArrayRef<mlir::Operation *> ops)
 {
     for (mlir::Operation *op : ops)