Skip to content

[mlir][linalg] Edit the yieldOutputs method's builder #73900

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

amirBish
Copy link
Contributor

Passing the RegionBuilder's builder to the yieldOutputs method, Instead of creating a new one within the yieldOutputs method. This change is being made since having memory leaks issues (when compiling with address sanitizer) when applying RewritePattern on LinalgStructuredOp (creating and erasing op within matchAndRewrite method would reproduce the issue).

@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir:linalg mlir:sparse Sparse compiler in MLIR mlir labels Nov 30, 2023
@llvmbot
Copy link
Member

llvmbot commented Nov 30, 2023

@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir-linalg

Author: Amir Bishara (amirBish)

Changes

Passing the RegionBuilder's builder to the yieldOutputs method, Instead of creating a new one within the yieldOutputs method. This change is being made since having memory leaks issues (when compiling with address sanitizer) when applying RewritePattern on LinalgStructuredOp (creating and erasing op within matchAndRewrite method would reproduce the issue).


Full diff: https://github.com/llvm/llvm-project/pull/73900.diff

3 Files Affected:

  • (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+2-2)
  • (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp (+1)
  • (modified) mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp (+1-1)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 58af9995548e939..11ebbabe9773787 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -480,8 +480,8 @@ class RegionBuilderHelper {
     llvm_unreachable("unsupported type conversion function");
   }
 
-  void yieldOutputs(ValueRange values) {
-    OpBuilder builder = getBuilder();
+  void yieldOutputs(OpBuilder builder, ValueRange values) {
+    builder.setInsertionPointToEnd(&block);
     Location loc = builder.getUnknownLoc();
     builder.create<YieldOp>(loc, values);
   }
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
index 375e10f9068e43b..e14efd8def0920d 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
@@ -260,6 +260,7 @@ struct SparseTensorCodegenPass
     // The following operations and dialects may be introduced by the
     // codegen rules, and are therefore marked as legal.
     target.addLegalOp<linalg::FillOp>();
+    target.addLegalOp<linalg::YieldOp>();
     target.addLegalDialect<
         arith::ArithDialect, bufferization::BufferizationDialect,
         complex::ComplexDialect, memref::MemRefDialect, scf::SCFDialect>();
diff --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
index fb3c9d48f9a9821..cd8398b973a5677 100644
--- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
+++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
@@ -1012,7 +1012,7 @@ void {0}::regionBuilder(ImplicitLocOpBuilder &b,
   SmallVector<Value> yields;
   {2}
   {3}
-  helper.yieldOutputs(yields);
+  helper.yieldOutputs(b, yields);
 }
 )FMT";
     auto &args = opConfig.structuredOp->args;

@amirBish
Copy link
Contributor Author

A reproduction for the stack trace that I've encountered
(the stack trace from an invalid lit test which uses -verify-diagnostics and runs a pass which applies a RewritePattern on linalg::CopyOp)

=================================================================
==1892478==ERROR: LeakSanitizer: detected memory leaks

Direct leak of 112 byte(s) in 1 object(s) allocated from:
    #0 0x7f70a59329f0 in __interceptor_malloc ../../../../libsanitizer/asan/asan_malloc_linux.cc:62
    #1 0xecda8db in mlir::Operation::create(mlir::Location, mlir::OperationName, mlir::TypeRange, mlir::ValueRange, mlir::DictionaryAttr, mlir::OpaqueProperties, mlir::BlockRange, unsigned int)
    #2 0xecdafaa in mlir::Operation::create(mlir::Location, mlir::OperationName, mlir::TypeRange, mlir::ValueRange, mlir::NamedAttrList&&, mlir::OpaqueProperties, mlir::BlockRange, unsigned int)
    #3 0xecdb199 in mlir::Operation::create(mlir::Location, mlir::OperationName, mlir::TypeRange, mlir::ValueRange, mlir::NamedAttrList&&, mlir::OpaqueProperties, mlir::BlockRange, mlir::RegionRange) 
    #4 0xecdc3fb in mlir::Operation::create(mlir::OperationState const&)
    #5 0xe8c7b6f in mlir::OpBuilder::create(mlir::OperationState const&) 
    #6 0x9376dd3 in (anonymous namespace)::RegionBuilderHelper::yieldOutputs(mlir::ValueRange)
    #7 0x9379faf in mlir::linalg::CopyOp::regionBuilder(mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>)
    #8 0x8c95fbb in std::_Function_handler<void (mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>), void (*)(mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>)>::_M_invoke(std::_Any_data const&, mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>&&) 
    #9 0x934a329 in void llvm::function_ref<void (mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>)>::callback_fn<std::function<void (mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>)> >(long, mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>)
    #10 0x9a19ee4 in fillStructuredOpRegion(mlir::OpBuilder&, mlir::Region&, mlir::TypeRange, mlir::TypeRange, llvm::ArrayRef<mlir::NamedAttribute>, llvm::function_ref<void (mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>)>) 
    #11 0x9a1bf2e in buildStructuredOp(mlir::OpBuilder&, mlir::OperationState&, std::optional<mlir::TypeRange>, mlir::ValueRange, mlir::ValueRange, llvm::ArrayRef<mlir::NamedAttribute>, llvm::function_ref<void (mlir::ImplicitLocOpBuilder&, mlir::Block&, llvm::ArrayRef<mlir::NamedAttribute>)>) 
    #12 0x9a2fb72 in mlir::linalg::CopyOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::ValueRange, mlir::ValueRange, llvm::ArrayRef<mlir::NamedAttribute>)
    #13 0x951fad in MemrefImportExportToAsyncRule::matchAndRewrite(mlir::linalg::CopyOp, mlir::PatternRewriter&) const::{lambda(mlir::OpBuilder&, mlir::Location, mlir::ValueRange)#1}::operator()(mlir::OpBuilder&, mlir::Location, mlir::ValueRange) const
    #14 0x9540e9 in void llvm::function_ref<void (mlir::OpBuilder&, mlir::Location, mlir::ValueRange)>::callback_fn<MemrefImportExportToAsyncRule::matchAndRewrite(mlir::linalg::CopyOp, mlir::PatternRewriter&) const::{lambda(mlir::OpBuilder&, mlir::Location, mlir::ValueRange)#1}>(long, mlir::OpBuilder&, mlir::Location, mlir::ValueRange)
    #15 0x3daf461 in mlir::async::ExecuteOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::TypeRange, mlir::ValueRange, mlir::ValueRange, llvm::function_ref<void (mlir::OpBuilder&, mlir::Location, mlir::ValueRange)>)
    #16 0x94f561 in MemrefImportExportToAsyncRule::matchAndRewrite(mlir::linalg::CopyOp, mlir::PatternRewriter&) const
    #17 0x94dfb0 in mlir::detail::OpOrInterfaceRewritePatternBase<mlir::linalg::CopyOp>::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const
    #18 0xabba210 in void llvm::function_ref<void ()>::callback_fn<mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>)::{lambda()#1}>(long)
    #19 0xabb8de4 in mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>)
    #20 0xa9a22b4 in (anonymous namespace)::OperationLegalizer::legalize(mlir::Operation*, mlir::ConversionPatternRewriter&)
    #21 0xa9a5046 in (anonymous namespace)::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation*>, llvm::function_ref<void (mlir::Diagnostic&)>)
    #22 0xa9b9231 in mlir::applyPartialConversion(llvm::ArrayRef<mlir::Operation*>, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, llvm::DenseSet<mlir::Operation*, llvm::DenseMapInfo<mlir::Operation*, void> >*)
    #23 0xa9b94fa in mlir::applyPartialConversion(mlir::Operation*, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, llvm::DenseSet<mlir::Operation*, llvm::DenseMapInfo<mlir::Operation*, void> >*)
    #24 0x95a867 in (anonymous namespace)::SomePass::runOnOperation()
    #25 0xc5970b9 in void llvm::function_ref<void ()>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::{lambda()#2}>(long)

SUMMARY: AddressSanitizer: 112 byte(s) leaked in 1 allocation(s).

void yieldOutputs(ValueRange values) {
OpBuilder builder = getBuilder();
void yieldOutputs(OpBuilder builder, ValueRange values) {
builder.setInsertionPointToEnd(&block);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the difference here exactly with before? Looking at the implementation of getBuilder() I'm not spotting it immediately.

Copy link
Contributor Author

@amirBish amirBish Nov 30, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the previous one the builder is being constructed from the block's argument context block.getArgument(0).getContext() , here is not the case It is being passed as is.
when using the rewriter of the MatchAndRewrite pattern we should use the rewriter which has been passed and not creating a new builder. That what has made the memory leakage.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ping

@@ -260,6 +260,7 @@ struct SparseTensorCodegenPass
// The following operations and dialects may be introduced by the
// codegen rules, and are therefore marked as legal.
target.addLegalOp<linalg::FillOp>();
target.addLegalOp<linalg::YieldOp>();
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is a test failing if you don't add that?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yep, without this legalizing YieldOp, the RewritePattern won't be applied (since YieldOp it's not legal).

@amirBish
Copy link
Contributor Author

amirBish commented Dec 8, 2023

@MaheshRavishankar are you familiar with these linalg methods? may you also please take a look on this change?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
mlir:core MLIR Core Infrastructure mlir:linalg mlir:sparse Sparse compiler in MLIR mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants