[MLIR] Generic 'malloc', 'aligned_alloc' and 'free' functions

[mscuttari]

Summary

When converted to the LLVM dialect, the memref.alloc and memref.free operations were generating calls to hardcoded 'malloc' and 'free' functions. This didn't leave any freedom to users to provide their custom implementation. Those operations now convert into calls to '_mlir_alloc' and '_mlir_free' functions, which have also been implemented into the runtime support library as wrappers to 'malloc' and 'free'. The same has been done for the 'aligned_alloc' function.

[lxsameer]

It would be nice to have this feature documented in the official docs.

[mscuttari]

In D128791#3618161, @lxsameer wrote:

It would be nice to have this feature documented in the official docs.

You're right, I will proceed in doing that.

[mscuttari]

When using JIT compilation, the _mlir_alloc and _mlir_free were not solved due to the missing runtime library. A new command line option has been added to get it loaded, and the tests have been updated accordingly.

[mscuttari]

Fix code formatting

[mscuttari]

In D128791#3618161, @lxsameer wrote:

It would be nice to have this feature documented in the official docs.

Well, I see that the official MLIR documentation is on another repo (mlir-www), so I think it should be better to patch it only when this modification lands. To be honest I don't know if I've to go through Phabricator to patch mlir-www, I've never done that.

[mscuttari]

Fix formatting

[myhsu]

I think you accidentally override the main patch, I can only see changes in the example files

[mscuttari]

Patch fix.

[mscuttari]

In D128791#3620465, @myhsu wrote:

I think you accidentally override the main patch, I can only see changes in the example files

First time using Phabricator, sorry. I made multiple commits to fix various problems, and each time uploaded a patch made with

git show HEAD -U999999 > mypatch.patch

Seems like it's not the way to go though. Please check if now it's ok, I've squashed all the changes into a single commit

[mscuttari]

Fix formatting

[myhsu]

In D128791#3621491, @mscuttari wrote:

In D128791#3620465, @myhsu wrote:

I think you accidentally override the main patch, I can only see changes in the example files

First time using Phabricator, sorry. I made multiple commits to fix various problems, and each time uploaded a patch made with

Right, we are using a different model than, say, GitHub PR. One of the reasons being that LLVM requires every commit to be buildable and able to pass all the tests, so one can't just stack new commits to update their changes. We, however, encourage people to split big changes into multiple patches, but of course each of them still needs to be buildable and passing all tests.

git show HEAD -U999999 > mypatch.patch

Seems like it's not the way to go though. Please check if now it's ok, I've squashed all the changes into a single commit

Usually I use git rebase to amend changes to a commit.

[myhsu]

Can you put some documentations (at mlir/docs) regarding this change? Otherwise LGTM.

mlir/examples/toy/Ch6/toyc.cpp

250 ↗

(On Diff #441305)

format: remove braces. https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

mlir/examples/toy/Ch7/toyc.cpp

251 ↗

(On Diff #441305)

ditto

[mscuttari]

Documentation updated & code style improvements

[mscuttari]

In D128791#3622897, @myhsu wrote:

Can you put some documentations (at mlir/docs) regarding this change? Otherwise LGTM.

I added a few lines into the Toy example documentation, please let me know if it's ok.

[mscuttari]

Ping. Unfortunately I do not have commit permissions so I need someone to commit the changes

[ftynse]

I am not very happy about exposing the complexity of shared libs and runnerutils to the tutorial. Can we just have a simple post-hoc transform (pass or just some helper code in the to JIT that goes through the symbol table of the module and renames _mlir_alloc back to malloc for the sake of simplicity. We can still mention that this is happening in the documentation and why, but it really looks desirable to avoid the complexity of shared library loading in the tutorial.

mlir/docs/Tutorials/Toy/Ch-6.md

174–176 ↗

(On Diff #441450)

Nit: please wrap at 80 cols.

[mscuttari]

Toy example: pass to rename '_mlir_alloc' into 'malloc' and '_mlir_free' into 'free'

[mscuttari]

In D128791#3638442, @ftynse wrote:

I am not very happy about exposing the complexity of shared libs and runnerutils to the tutorial. Can we just have a simple post-hoc transform (pass or just some helper code in the to JIT that goes through the symbol table of the module and renames _mlir_alloc back to malloc for the sake of simplicity. We can still mention that this is happening in the documentation and why, but it really looks desirable to avoid the complexity of shared library loading in the tutorial.

I've introduced a small transformation pass into the Toy example, and fixed the documentation of chapter 6 accordingly.

[mscuttari]

Code formatting

[mscuttari]

MSVC's implementation of the standard library doesn't provide the aligned_alloc function, despite it being part of the standard. Instead, it provides an _aligned_malloc function. Moreover, MSVC also requires the pointers obtained through _aligned_malloc to be deallocated through _aligned_free, and not free.

[mscuttari]

Ping

[ftynse]

Sorry for the delay.

[mscuttari]

In D128791#3659984, @ftynse wrote:

Sorry for the delay.

No problem. However I will need someone to commit the patch, I do not have the rights to do it by myself

[ftynse]

Could you please provide your name and email as they should be written in the git commit?

[mscuttari]

Yes sure:

Michele Scuttari
[e-mail]

Thanks

[mehdi_amini]

I had to revert: all the integration tests are broken. Try to reconfigure your build dir with -DMLIR_INCLUDE_INTEGRATION_TESTS=ON

[silvas]

FYI this caused us issues in Torch-MLIR as well https://github.com/llvm/torch-mlir/pull/1078#issuecomment-1188337337

I'm a little worried about the direction on this change -- LLVM already permits malloc/free to be treated specially, so I don't see why we need to do this renaming at the MLIR level. I think the fact that we needed a special pass to keep the Toy example running indicates that this is going in the wrong direction (we should not be moving towards making tutorials more complicated / less "just works" unless there is a very good reason). It seems like the better direction would be to provide a pass that replaces malloc and free with _mlir_malloc/free which users can hook into the pipeline if they desire the renaming -- this pass could even be an LLVM IR pass.

[mscuttari]

In D128791#3661185, @silvas wrote:

FYI this caused us issues in Torch-MLIR as well https://github.com/llvm/torch-mlir/pull/1078#issuecomment-1188337337

I'm a little worried about the direction on this change -- LLVM already permits malloc/free to be treated specially, so I don't see why we need to do this renaming at the MLIR level. I think the fact that we needed a special pass to keep the Toy example running indicates that this is going in the wrong direction (we should not be moving towards making tutorials more complicated / less "just works" unless there is a very good reason). It seems like the better direction would be to provide a pass that replaces malloc and free with _mlir_malloc/free which users can hook into the pipeline if they desire the renaming -- this pass could even be an LLVM IR pass.

I'm sorry for the issues, I didn't know about the MLIR_INCLUDE_INTEGRATION_TESTS option and neither I could find it documented anywhere, so I thought that passing the Phabricator builds and the mlir-check target tests was enough to ensure I didn't break anything.
I can try implementing the pass you mentioned, which basically is the opposite of what is done inside the Toy example. At that point, however, I don't know if it has a reason to exist at all and if it should instead left to be implemented by the user.
The rationale behind this change was to avoid sticking to malloc and free, which seem like first-class citizens if one ignores that the standard library is automatically linked (same story for the printf function, which is the Toy example). Also, there are some caveats such as the aligned_alloc function, which MSVC's implementation does not provide (yet is used by MLIR in some places), and also imply different deallocation calls. For this last issue, however, we can just keep the related part of the patch, if you think it may be useful (at that point, however, there would be a naming inconsistency which I don't know how much we would like).
Anyway thank you for your share, I've been using MLIR as a user for the last two years but I'm quite new in delivering changes upstream, so I would really apreciate more opinions on this topic so that we can decide the path to take (in which doing nothing is absolutely one of them, if you believe so).

[mehdi_amini]

In D128791#3661671, @mscuttari wrote:

I'm sorry for the issues, I didn't know about the MLIR_INCLUDE_INTEGRATION_TESTS option and neither I could find it documented anywhere, so I thought that passing the Phabricator builds and the mlir-check target tests was enough to ensure I didn't break anything.

Don't worry: it is hard to not break "something".
For example we also have a bot running code on GPUs which you may not be able to test locally if you don't have a Nvidia GPU. We also have some bots that are big endian which again you may have a hard time running locally.

[mehdi_amini]

One thing though: be on the lookout in the next hour after pushing a change, you may get email notifications from bots that are broken. You'll need to determine if it is related to your patch (sometimes multiple changes get landed closely together and the bot would test them together) and quickly fix forward or revert.

[ftynse]

There was a forum post and nobody really raised any objection to the direction there -- https://discourse.llvm.org/t/llvm-dialect-replacing-malloc-and-free-with-custom-functions/63481. FWIW, the original lowering of memref allocs did use custom functions -- https://github.com/llvm/llvm-project/commit/90d1b6b5f25e66059be5be1f9badcbc5a37c356b. It was later silently changed in https://github.com/llvm/llvm-project/commit/e9493cf14deec4198a3620d734f03e7e143f91d6, with as motivation being able to run a JIT without a support library. IMO, we (well, I as the commit author) took a shortcut there. A custom function lets us intercept and customize specifically the allocations coming from the memref dialect and ignore everything else. A pass that would renamed malloc to _mlir_alloc will also do so for legitimate user calls to malloc, and it may not be what we want. Maybe we can have a binary switch the lowering pass between emitting _mlir_alloc and malloc, defaulting to _mlir_alloc as the more scalable approach. Same as we have the non-default bare pointer calling convention that the user may opt into if their case is simple enough. This is slightly better than having to inject function names as strings through pass options.

[mscuttari]

In D128791#3662012, @mehdi_amini wrote:

One thing though: be on the lookout in the next hour after pushing a change, you may get email notifications from bots that are broken. You'll need to determine if it is related to your patch (sometimes multiple changes get landed closely together and the bot would test them together) and quickly fix forward or revert.

Probably I didn't receive any notification because I asked Alex to commit it. The next times I will pay attention to the Actions list on github and see if anything strange happens.
Just to confirm: the bots leveraged by Phabricator do check if the build process succeeds, but they do not run the tests? Because I was getting a green check when I was uploading my patches, and that made me believe everything was fine.

In D128791#3662641, @ftynse wrote:

Maybe we can have a binary switch the lowering pass between emitting _mlir_alloc and malloc, defaulting to _mlir_alloc as the more scalable approach. Same as we have the non-default bare pointer calling convention that the user may opt into if their case is simple enough. This is slightly better than having to inject function names as strings through pass options.

I will look into this asap.

[ftynse]

In D128791#3662711, @mscuttari wrote:

In D128791#3662012, @mehdi_amini wrote:

Probably I didn't receive any notification because I asked Alex to commit it. The next times I will pay attention to the Actions list on github and see if anything strange happens.
Just to confirm: the bots leveraged by Phabricator do check if the build process succeeds, but they do not run the tests? Because I was getting a green check when I was uploading my patches, and that made me believe everything was fine.

I did not receive it either. The email gets sent to the address associated with the commit, the one you gave is @users.noreply.github.com so the email was most likely bounced by GitHub. IIRC, there was a way to get a "reply" email from github that gets forwarded somewhere.
Some notifications don't go through GitHub, so it's better to check the email...

In D128791#3662641, @ftynse wrote:

Maybe we can have a binary switch the lowering pass between emitting _mlir_alloc and malloc, defaulting to _mlir_alloc as the more scalable approach. Same as we have the non-default bare pointer calling convention that the user may opt into if their case is simple enough. This is slightly better than having to inject function names as strings through pass options.

I will look into this asap.

Thanks!

[mscuttari]

In D128791#3662728, @ftynse wrote:

I did not receive it either. The email gets sent to the address associated with the commit, the one you gave is @users.noreply.github.com so the email was most likely bounced by GitHub. IIRC, there was a way to get a "reply" email from github that gets forwarded somewhere.
Some notifications don't go through GitHub, so it's better to check the email...

Oh well you're right. I usually get notifications in case of CI failures on my project but, being the bots external, the email surely got reject by Github. Next time I will provide you my personal email (maybe privately) so I can get notified in case of problems and still have it linked to the account.

[ftynse]

If you upload patches using Arcanist, the email in your commit will get included in the diff metadata and I will be able to use it directly.

[mehdi_amini]

In D128791#3662711, @mscuttari wrote:

Just to confirm: the bots leveraged by Phabricator do check if the build process succeeds, but they do not run the tests? Because I was getting a green check when I was uploading my patches, and that made me believe everything was fine.

It runs the usual tests, but likely not the integration one? (that is it does not have the CMake option I mentioned to you).

[silvas]

In D128791#3662641, @ftynse wrote:

There was a forum post and nobody really raised any objection to the direction there -- https://discourse.llvm.org/t/llvm-dialect-replacing-malloc-and-free-with-custom-functions/63481. FWIW, the original lowering of memref allocs did use custom functions -- https://github.com/llvm/llvm-project/commit/90d1b6b5f25e66059be5be1f9badcbc5a37c356b. It was later silently changed in https://github.com/llvm/llvm-project/commit/e9493cf14deec4198a3620d734f03e7e143f91d6, with as motivation being able to run a JIT without a support library. IMO, we (well, I as the commit author) took a shortcut there. A custom function lets us intercept and customize specifically the allocations coming from the memref dialect and ignore everything else. A pass that would renamed malloc to _mlir_alloc will also do so for legitimate user calls to malloc, and it may not be what we want. Maybe we can have a binary switch the lowering pass between emitting _mlir_alloc and malloc, defaulting to _mlir_alloc as the more scalable approach. Same as we have the non-default bare pointer calling convention that the user may opt into if their case is simple enough. This is slightly better than having to inject function names as strings through pass options.

In that case I think we should call it _mlir_memref_dialect_alloc and not _mlir_alloc.

I still think that things should "just work" and customizations like this should be graceful improvements rather than up-front cognitive burden for all users. So having this as a flag (default no change from current behavior) and very specific to the memref dialect to LLVM conversion would be ideal in my mind.

[mscuttari]

Option inside MemRef -> LLVM conversion pass

'use-generic-function' option inside the MemRef -> LLVM conversion pass to enable
the usage of generic allocation / deallocation functions. When set to true,
'_mlir_alloc' is used instead of 'malloc', '_mlir_aligned_alloc' instead of
'aligned_alloc' and '_mlir_free' instead of 'free'. The option defaults to false.

[mscuttari]

Well I'm sorry for the stupid question but I'm quite confused. I switched to Arcanist but things seem to be messed up.
I made the changes we agreed on, committed them, reopened the revision and finally run arc diff --update D128791.
Still the changes seem to be the older ones, and I see no trace of the last diff. Can anyone please explain what I'm doing wrong?

EDIT: should be ok now, but still a bit confused about the overall flow, will see if it gets better with future patches

[mscuttari]

Option inside MemRef -> LLVM conversion pass

'use-generic-function' option inside the MemRef -> LLVM conversion pass to enable the usage of generic allocation / deallocation functions. When set to true, '_mlir_alloc' is used instead of 'malloc', '_mlir_aligned_alloc' instead of 'aligned_alloc' and '_mlir_free' instead of 'free'. The option defaults to false.

[mscuttari]

Fix formatting

[ftynse]

I made the changes we agreed on, committed them, reopened the revision and finally run arc diff --update D128791.

You need to specify what you are diff'ing against, e.g., the previous commit HEAD^.

[mscuttari]

In D128791#3675981, @ftynse wrote:

I made the changes we agreed on, committed them, reopened the revision and finally run arc diff --update D128791.

You need to specify what you are diff'ing against, e.g., the previous commit HEAD^.

Oh right, thanks. I completely forgot as I was previously doing the diff manually.
If the patch is good for @silvas (which I think it is, as the default behaviour has been preserved), then I would kindly ask you commit the changes on the repo. You should also now be able to see an email that is different from the github one which I previously wrote. Please tell me if this is not the case, or I will not be able to see problems with the build-bots, if any should ever arise.

[ftynse]

You can always check the build status manually: https://lab.llvm.org/buildbot/#/changes/64333

[silvas]

I would recommend that we rename this to _mlir_memref_to_llvm_alloc, to differentiate it from other situations where MLIR might want a custom allocation function (memref to LLVM is not "MLIR" as whole)

[mscuttari]

In D128791#3677076, @silvas wrote:

I would recommend that we rename this to _mlir_memref_to_llvm_alloc, to differentiate it from other situations where MLIR might want a custom allocation function (memref to LLVM is not "MLIR" as whole)

Alright, I'm sorry but I forgot that, I see now that you already mentioned it in your previous comment.
Newbie question about the Phabricator's workflow: should I open a new revision or reopen this one?

[silvas]

You can open a new one -- thanks!

mlir/docs/TargetLLVMIR.md

 }
 ```
 
 #### Bare Pointer Calling Convention For Unranked MemRef
 
 The "bare pointer" calling convention does not support unranked memrefs as their
 shape cannot be known at compile time.
 
+### Generic alloction and deallocation functions
+
+When converting the Memref dialect, allocations and deallocations are converted
+into calls to `malloc` (`aligned_alloc` if aligned allocations are requested)
+and `free`. However, it is possible to convert them to more generic functions
+which can be implemented by a runtime library, thus allowing custom allocation
+strategies or runtime profiling. When the conversion pass is  instructed to
+perform such operation, the names of the calles are `_mlir_alloc`,
+`_mlir_aligned_alloc` and `_mlir_free`. Their signatures are the same of
+`malloc`, `aligned_alloc` and `free`.
+
 ### C-compatible wrapper emission
 
 In practical cases, it may be desirable to have externally-facing functions with
 a single attribute corresponding to a MemRef argument. When interfacing with
 LLVM IR produced from C, the code needs to respect the corresponding calling
 convention. The conversion to the LLVM dialect provides an option to generate
 wrapper functions that take memref descriptors as pointers-to-struct compatible
 with data types produced by Clang when compiling C sources. The generation of

mlir/include/mlir/Conversion/LLVMCommon/LoweringOptions.h

enum class AllocLowering {
 
     /// Do not lower heap allocations. Users must provide their own patterns for
     /// AllocOp and DeallocOp lowering.
     None
   };
 
   AllocLowering allocLowering = AllocLowering::Malloc;
 
+  bool useGenericFunctions = false;
+
   /// The data layout of the module to produce. This must be consistent with the
   /// data layout used in the upper levels of the lowering pipeline.
   // TODO: this should be replaced by MLIR data layout when one exists.
   llvm::DataLayout dataLayout = llvm::DataLayout("");
 
   /// Set the index bitwidth to the given value.
   void overrideIndexBitwidth(unsigned bitwidth) {
     assert(bitwidth != kDeriveIndexBitwidthFromDataLayout &&

mlir/include/mlir/Conversion/Passes.td

def ConvertMemRefToLLVM : Pass<"convert-memref-to-llvm", "ModuleOp"> {
   let constructor = "mlir::createMemRefToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"useAlignedAlloc", "use-aligned-alloc", "bool", /*default=*/"false",
            "Use aligned_alloc in place of malloc for heap allocations">,
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
+    Option<"useGenericFunctions", "use-generic-functions",
+           "bool",
+           /*default=*/"false",
+           "Use generic allocation and deallocation functions instead of the "
+           "classic 'malloc', 'aligned_alloc' and 'free' functions">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefToSPIRV
 //===----------------------------------------------------------------------===//
 
 def ConvertMemRefToSPIRV : Pass<"convert-memref-to-spirv", "ModuleOp"> {

mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h

 LLVM::LLVMFuncOp lookupOrCreatePrintOpenFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintCloseFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintCommaFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintNewlineFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreateMallocFn(ModuleOp moduleOp, Type indexType);
 LLVM::LLVMFuncOp lookupOrCreateAlignedAllocFn(ModuleOp moduleOp,
                                               Type indexType);
 LLVM::LLVMFuncOp lookupOrCreateFreeFn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreateGenericAllocFn(ModuleOp moduleOp,
+                                              Type indexType);
+LLVM::LLVMFuncOp lookupOrCreateGenericAlignedAllocFn(ModuleOp moduleOp,
+                                                     Type indexType);
+LLVM::LLVMFuncOp lookupOrCreateGenericFreeFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreateMemRefCopyFn(ModuleOp moduleOp, Type indexType,
                                             Type unrankedDescriptorType);
 
 /// Create a FuncOp with signature `resultType`(`paramTypes`)` and name `name`.
 LLVM::LLVMFuncOp lookupOrCreateFn(ModuleOp moduleOp, StringRef name,
                                   ArrayRef<Type> paramTypes = {},
                                   Type resultType = {});
 

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

bool isStaticStrideOrOffset(int64_t strideOrOffset) {
   return !ShapedType::isDynamicStrideOrOffset(strideOrOffset);
 }
 
 struct AllocOpLowering : public AllocLikeOpLLVMLowering {
   AllocOpLowering(LLVMTypeConverter &converter)
       : AllocLikeOpLLVMLowering(memref::AllocOp::getOperationName(),
                                 converter) {}
 
+  LLVM::LLVMFuncOp getAllocFn(ModuleOp module) const {
+    bool useGenericFn = getTypeConverter()->getOptions().useGenericFunctions;
+
+    if (useGenericFn)
+      return LLVM::lookupOrCreateGenericAllocFn(module, getIndexType());
+
+    return LLVM::lookupOrCreateMallocFn(module, getIndexType());
+  }
+
   std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                           Location loc, Value sizeBytes,
                                           Operation *op) const override {
     // Heap allocations.
     memref::AllocOp allocOp = cast<memref::AllocOp>(op);
     MemRefType memRefType = allocOp.getType();
 
     Value alignment;
std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
     if (alignment) {
       // Adjust the allocation size to consider alignment.
       sizeBytes = rewriter.create<LLVM::AddOp>(loc, sizeBytes, alignment);
     }
 
     // Allocate the underlying buffer and store a pointer to it in the MemRef
     // descriptor.
     Type elementPtrType = this->getElementPtrType(memRefType);
-    auto allocFuncOp = LLVM::lookupOrCreateMallocFn(
-        allocOp->getParentOfType<ModuleOp>(), getIndexType());
+    auto allocFuncOp = getAllocFn(allocOp->getParentOfType<ModuleOp>());
     auto results = createLLVMCall(rewriter, loc, allocFuncOp, {sizeBytes},
                                   getVoidPtrType());
     Value allocatedPtr =
         rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, results[0]);
 
     Value alignedPtr = allocatedPtr;
     if (alignment) {
       // Compute the aligned type pointer.
int64_t getAllocationAlignment(memref::AllocOp allocOp) const {
     // Whenever we don't have alignment set, we will use an alignment
     // consistent with the element type; since the allocation size has to be a
     // power of two, we will bump to the next power of two if it already isn't.
     auto eltSizeBytes = getMemRefEltSizeInBytes(allocOp.getType(), allocOp);
     return std::max(kMinAlignedAllocAlignment,
                     llvm::PowerOf2Ceil(eltSizeBytes));
   }
 
+  LLVM::LLVMFuncOp getAllocFn(ModuleOp module) const {
+    bool useGenericFn = getTypeConverter()->getOptions().useGenericFunctions;
+
+    if (useGenericFn)
+      return LLVM::lookupOrCreateGenericAlignedAllocFn(module, getIndexType());
+
+    return LLVM::lookupOrCreateAlignedAllocFn(module, getIndexType());
+  }
+
   std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                           Location loc, Value sizeBytes,
                                           Operation *op) const override {
     // Heap allocations.
     memref::AllocOp allocOp = cast<memref::AllocOp>(op);
     MemRefType memRefType = allocOp.getType();
     int64_t alignment = getAllocationAlignment(allocOp);
     Value allocAlignment = createIndexConstant(rewriter, loc, alignment);
 
     // aligned_alloc requires size to be a multiple of alignment; we will pad
     // the size to the next multiple if necessary.
     if (!isMemRefSizeMultipleOf(memRefType, alignment, op))
       sizeBytes = createAligned(rewriter, loc, sizeBytes, allocAlignment);
 
     Type elementPtrType = this->getElementPtrType(memRefType);
-    auto allocFuncOp = LLVM::lookupOrCreateAlignedAllocFn(
-        allocOp->getParentOfType<ModuleOp>(), getIndexType());
+    auto allocFuncOp = getAllocFn(allocOp->getParentOfType<ModuleOp>());
     auto results =
         createLLVMCall(rewriter, loc, allocFuncOp, {allocAlignment, sizeBytes},
                        getVoidPtrType());
     Value allocatedPtr =
         rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, results[0]);
 
     return std::make_tuple(allocatedPtr, allocatedPtr);
   }

 // The memref descriptor being an SSA value, there is no need to clean it up
 // in any way.
 struct DeallocOpLowering : public ConvertOpToLLVMPattern<memref::DeallocOp> {
   using ConvertOpToLLVMPattern<memref::DeallocOp>::ConvertOpToLLVMPattern;
 
   explicit DeallocOpLowering(LLVMTypeConverter &converter)
       : ConvertOpToLLVMPattern<memref::DeallocOp>(converter) {}
 
+  LLVM::LLVMFuncOp getFreeFn(ModuleOp module) const {
+    bool useGenericFn = getTypeConverter()->getOptions().useGenericFunctions;
+
+    if (useGenericFn)
+      return LLVM::lookupOrCreateGenericFreeFn(module);
+
+    return LLVM::lookupOrCreateFreeFn(module);
+  }
+
   LogicalResult
   matchAndRewrite(memref::DeallocOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Insert the `free` declaration if it is not already present.
-    auto freeFunc = LLVM::lookupOrCreateFreeFn(op->getParentOfType<ModuleOp>());
+    auto freeFunc = getFreeFn(op->getParentOfType<ModuleOp>());
     MemRefDescriptor memref(adaptor.getMemref());
     Value casted = rewriter.create<LLVM::BitcastOp>(
         op.getLoc(), getVoidPtrType(),
         memref.allocatedPtr(rewriter, op.getLoc()));
     rewriter.replaceOpWithNewOp<LLVM::CallOp>(
         op, TypeRange(), SymbolRefAttr::get(freeFunc), casted);
     return success();
   }
struct MemRefToLLVMPass : public ConvertMemRefToLLVMBase<MemRefToLLVMPass> {
   void runOnOperation() override {
     Operation *op = getOperation();
     const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
     LowerToLLVMOptions options(&getContext(),
                                dataLayoutAnalysis.getAtOrAbove(op));
     options.allocLowering =
         (useAlignedAlloc ? LowerToLLVMOptions::AllocLowering::AlignedAlloc
                          : LowerToLLVMOptions::AllocLowering::Malloc);
+
+    options.useGenericFunctions = useGenericFunctions;
+
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
 
     LLVMTypeConverter typeConverter(&getContext(), options,
                                     &dataLayoutAnalysis);
     RewritePatternSet patterns(&getContext());
     populateMemRefToLLVMConversionPatterns(typeConverter, patterns);
     LLVMConversionTarget target(getContext());

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp

 static constexpr llvm::StringRef kPrintF64 = "printF64";
 static constexpr llvm::StringRef kPrintOpen = "printOpen";
 static constexpr llvm::StringRef kPrintClose = "printClose";
 static constexpr llvm::StringRef kPrintComma = "printComma";
 static constexpr llvm::StringRef kPrintNewline = "printNewline";
 static constexpr llvm::StringRef kMalloc = "malloc";
 static constexpr llvm::StringRef kAlignedAlloc = "aligned_alloc";
 static constexpr llvm::StringRef kFree = "free";
+static constexpr llvm::StringRef kGenericAlloc = "_mlir_alloc";
+static constexpr llvm::StringRef kGenericAlignedAlloc = "_mlir_aligned_alloc";
+static constexpr llvm::StringRef kGenericFree = "_mlir_free";
 static constexpr llvm::StringRef kMemRefCopy = "memrefCopy";
 
 /// Generic print function lookupOrCreate helper.
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateFn(ModuleOp moduleOp, StringRef name,
                                               ArrayRef<Type> paramTypes,
                                               Type resultType) {
   auto func = moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(name);
   if (func)

 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateFreeFn(ModuleOp moduleOp) {
   return LLVM::lookupOrCreateFn(
       moduleOp, kFree,
       LLVM::LLVMPointerType::get(IntegerType::get(moduleOp->getContext(), 8)),
       LLVM::LLVMVoidType::get(moduleOp->getContext()));
 }
 
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateGenericAllocFn(ModuleOp moduleOp,
+                                                          Type indexType) {
+  return LLVM::lookupOrCreateFn(
+      moduleOp, kGenericAlloc, indexType,
+      LLVM::LLVMPointerType::get(IntegerType::get(moduleOp->getContext(), 8)));
+}
+
+LLVM::LLVMFuncOp
+mlir::LLVM::lookupOrCreateGenericAlignedAllocFn(ModuleOp moduleOp,
+                                                Type indexType) {
+  return LLVM::lookupOrCreateFn(
+      moduleOp, kGenericAlignedAlloc, {indexType, indexType},
+      LLVM::LLVMPointerType::get(IntegerType::get(moduleOp->getContext(), 8)));
+}
+
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateGenericFreeFn(ModuleOp moduleOp) {
+  return LLVM::lookupOrCreateFn(
+      moduleOp, kGenericFree,
+      LLVM::LLVMPointerType::get(IntegerType::get(moduleOp->getContext(), 8)),
+      LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+
 LLVM::LLVMFuncOp
 mlir::LLVM::lookupOrCreateMemRefCopyFn(ModuleOp moduleOp, Type indexType,
                                        Type unrankedDescriptorType) {
   return LLVM::lookupOrCreateFn(
       moduleOp, kMemRefCopy,
       ArrayRef<Type>{indexType, unrankedDescriptorType, unrankedDescriptorType},
       LLVM::LLVMVoidType::get(moduleOp->getContext()));
 }

mlir/test/Conversion/MemRefToLLVM/generic-functions.mlir

+// RUN: mlir-opt -pass-pipeline="convert-memref-to-llvm{use-generic-functions=1}" -split-input-file %s \
+// RUN: | FileCheck %s --check-prefix="CHECK-NOTALIGNED"
+
+// RUN: mlir-opt -pass-pipeline="convert-memref-to-llvm{use-generic-functions=1 use-aligned-alloc=1}" -split-input-file %s \
+// RUN: | FileCheck %s --check-prefix="CHECK-ALIGNED"
+
+// CHECK-LABEL: func @alloc()
+func.func @zero_d_alloc() -> memref<f32> {
+// CHECK-NOTALIGNED: llvm.call @_mlir_alloc(%{{.*}}) : (i64) -> !llvm.ptr<i8>
+// CHECK-ALIGNED: llvm.call @_mlir_aligned_alloc(%{{.*}}, %{{.*}}) : (i64, i64) -> !llvm.ptr<i8>
+  %0 = memref.alloc() : memref<f32>
+  return %0 : memref<f32>
+}
+
+// -----
+
+// CHECK-LABEL: func @dealloc()
+func.func @dealloc(%arg0: memref<f32>) {
+// CHECK-NOTALIGNED: llvm.call @_mlir_free(%{{.*}}) : (!llvm.ptr<i8>) -> ()
+// CHECK-ALIGNED: llvm.call @_mlir_free(%{{.*}}) : (!llvm.ptr<i8>) -> ()
+  memref.dealloc %arg0 : memref<f32>
+  return
+}