[MLIR] Add support to use aligned_alloc to lower AllocOp from std to llvm

[bondhugula]

Summary

Support to recognize and deal with aligned_alloc was recently added to
LLVM's TLI/MemoryBuiltins and its various optimization passes. This
revision adds support for generation of aligned_alloc's when lowering
AllocOp from std to LLVM. Setting 'use-aligned_alloc=1' will lead to
aligned_alloc being used for all heap allocations. An alignment and size
that works with the constraints of aligned_alloc is chosen.

Using aligned_alloc is preferable to "using malloc and adjusting the
allocated pointer to align for indexing" because the pointer access
arithmetic done for the latter only makes it harder for LLVM passes to
deal with for analysis, optimization, attribute deduction, and rewrites.

Depends on D76602.

[mehdi_amini]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

50

Please don't introduce any new global command line option and thread them through the APIs.
We're looking to eliminate every global initializers under mlir/lib/... : https://bugs.llvm.org/show_bug.cgi?id=45437

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1397

Can't you just rely on the Optional result being None to signal that aligned alloc is not required?

1418

Prefer early return

1420

I wouldn't add a lambda that is only called once immediately after its definition.

1478

Nit: putting the comment above if makes it less surprising in terms of indentation. Otherwise, it feels like there is incorrect indentation because of elided braces and the actual statement does not belong to the`if`.

1582

I am a bit concerned about hardcoding the "typical" value. Can we make it parametric instead?

[bondhugula]

Thanks for the review.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1397

We still use aligned_alloc if the elt size is > 16 bytes (in spite of no alignment attribute). The code would otherwise crash since the loads/stores won't have alignment guaranteed (with malloc) but the they would be lowered by LLVM via aligned load/stores. (LLVM in the absence of alignment attributes on load/store would assume ABI alignment which for say vector<8xf32> elt types would be 256-bit boundaries).

1418

Thanks.

1420

Hmm... this is just for better readability - it gives a name / auto documents a code block without the need to outline it into a function or add an explicit comment. I've seen this as a standard practice.

1478

Sure, thanks.

1582

I had sort of a similar concern. But 16 bytes is pretty much what glibc malloc gives on nearly every system we have (on probably really old ones, it was perhaps 8 bytes). Did you want a pass flag and then letting 16 be the default - that would be too much plumbing (just like alignedAlloc). This is already a parameter of sorts.

[bondhugula]

Address review comments.

[ftynse]

mlir/include/mlir/Conversion/Passes.td

228

Nit: I suppose this option should be removed by the previous commit, instead of rewritten by this one.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1397

We still use aligned_alloc if the elt size is > 16 bytes

At which point, this function returns the alignment value rather than llvm::None.
The only situation where it returns a non-None value and sets useAlignedAlloc to false is for AllocaOp. I would consider refactoring this function to only work on AllocOp, and querying the AllocaOp allocation separately in the call place. This would make the API simpler and shorten the code. And you wouldn't need an extra boolean result.

(LLVM in the absence of alignment attributes on load/store would assume ABI alignment which for say vector<8xf32> elt types would be 256-bit boundaries).

I suppose we can just take the ABI alignment into account when lowering the alloc operation (we already do so for the alignment attribute). Not all platforms have aligned_alloc.

1420

This does not seem to be common practice in MLIR. FWIW, I find it less readable than just writing

int64_t constEltSizeBytes = 0;
if (auto vectorType = elementType.template dyn_cast<VectorType>())
  constEltSizeBytes =
      vectorType.getNumElements() *
      llvm::divideCeil(vectorType.getElementTypeBitWidth(), 8);
else
  constEltSizeBytes =
      llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);

// Use aligned_alloc if elt_size > malloc's alignment.
bool isMallocAlignmentSufficient = constEltSizeBytes > kMallocAlignment;
useAlignedAlloc = isMallocAlignmentSufficient;

Since you already have the block comment immediately above it anyway, and variables can have names just as well as lambdas. The lambda also mutates a global state that it captures by-reference, so the only effects of lambda are: (1) extra indentation; (2) extra LoC; and (3) extra concepts leading to cognitive overhead.

1464

I would just have

auto allocaOp = dyn_cast<AllocaOp>(op);
if (allocaOp) {
  allocatedBytePtr = nullptr;
  accessAlignment = nullptr;
  return rewriter.create<LLVM::AllocaOp>(
          loc, elementPtrType, cumulativeSize,
          allocaOp.alignment() ? allocaOp.alignemnt().getValue() : 0);
}

and sink the getAllocationAlignment below, all while making it handle only AllocOp.

1582

The world is not limited to glibc. MLIR should also work on other platforms, and you essentially shift the burden of the plumbing you didn't do to people debugging builds on those platforms. You can have one pass option that corresponds to malloc alignment and, if it is set to 0, treat it as "never use aligned_alloc".

[bondhugula]

Address some of the review comments.

[bondhugula]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1397

I would consider refactoring this function to only work on
AllocOp, and querying the AllocaOp allocation separately in the

Sure, sounds good.

I suppose we can just take the ABI alignment into account when
lowering the alloc operation (we already do so for the

Yes, we should do this with malloc as well. This can be a TODO and should be done in a separate revision (because it isn't to do with aligned_alloc but fixing malloc lowering).

1420

Hmm... the demarcation/isolation is important I feel. I'm fine with changing to the straightline style but out of curiosity and for future purposes as well, it'll be good to have a third person view here on coding style as far as such patterns go: @mehdi_amini - is there a guideline here?

1464

Sure.

1582

Sorry, I didn't quite understand. What should the pass options be and what should the behavior and the default behavior be?

[bondhugula]

Address most of the review comments.

[bondhugula]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1397

Not all platforms have aligned_alloc.

Btw, aligned_alloc is a C11 and C++11 standard, and it's available on Windows too (it's posix_memalign that isn't). And of course there are systems where there isn't any dynamic memory allocation at all.

[bondhugula]

Remove stale comment.

[nicolasvasilache]

It is my understanding this CL does not change previous behavior, is this accurate?

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

45

Same comment as in the parent revision, this is a silently breaking API change that will impact integrations.
I'd just use the struct that you'd have introduced in the previous revision and document that this is only using 2 fields.

66–67

same discussion re silently breaking API changes

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1408

Alignment needed is really a target-specific thing isn't it?
I would have expected something that looks at DataLayout and that brings in the can of worms we have been punting on (I understand once flang is in MLIR core we will want to reopen it).

Can this part be dropped from this revision, esp in light of @ftynse's comments?

I also have some fun micro ARM targets I will need to test some of this on, the smaller the baked in assumptions the better.

1420

I am generally a fan of such style (esp. when mixed with functional combinators), so I'd vote for +1 it when it makes sense.

[mehdi_amini]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1420

@bondhugula this seems too detailed to have a guideline :)
I wouldn't say it is "common", but probably not unheard of?
I have been doing this myself but in general not calling it right after, rather to outline a block of code outside of a loop to make the loop shorter and easier to read, or similar situation (getting large boilerplate out of the way and "naming it").

@rriddle ?

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1420

Generally, I would be strongly opposed to defining style guidelines based on a single use of a construct in a single diff, where only a small subset of contributors could participate (and before you can object that review history is public, are you reading all comments on all diffs?). I would be also opposed to having to define additional rules until we have to. I am not generally opposed to helper lambdas, I just don't see any benefit from this specific one, only drawbacks. And a lambda that the entire environment by reference is not exactly my definition of isolation.

1582

Normally, you would have two pass options (and a configuration struct for the constructors like Nicolas proposed in another patch to decrease the amount of churn in pass constructor APIs): -use-aligned-alloc and -assume-malloc-alignment. If you don't want two separate options, you could get away with one -use-aligned-alloc-and-assume-malloc-alignment (did not think about a better name). If it is set to zero (default), the conversion doesn't use aligned_alloc at all. If it is set to non-zero, the conversion uses aligned_alloc and treats the option value as malloc alignment in order to also use aligned_alloc in relevant cases.

[bondhugula]

LowerToLLVM struct for options.

[bondhugula]

@nicolasvasilache wrote:

It is my understanding this CL does not change previous behavior, is this accurate?

It actually changes it in one case: if the alloc op didn't have alignment specified *and* its elt type was larger than 16 bytes, this revision would make it use aligned_alloc. I can change this and let it continue to use malloc, but note that the malloc in such cases would have generated code that would crash. So this is only fixing things for the previous behavior it's changing. And we should also change malloc path to do the alignment to elt type boundaries for > 16 bytes by default.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1420

As Mehdi now confirms, this is too detailed to have a style guideline. The lambda demarcates the start and end of the thing it's naming/auto-documenting - you don't get it from code comments alone and I see it better for readability. Given @ntv's and @mehdi_amini's comments, I'm now strongly inclined to retain it.

[bondhugula]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1582

The configuration struct is now done (in the parent revision). How about just keeping it simple with -use-aligned-alloc and not changing previous/existing behavior when -use-aligned-alloc is not provided? This revision is not about tinkering with malloc alignment handling.

Update - PTAL. Thanks for all the feedback.

[bondhugula]

Missed change.

[bondhugula]

Remove debug code.

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1397

Yes, I know it's C11 (but C++17, most changes in C11 did not make it to C++11), but unfortunately, I don't think we can assume all target systems have C11, let alone C++17.

1420

The lambda demarcates the start and end of the thing it's naming/auto-documenting - you don't get it from code comments alone and I see it better for readability.

Well, you have a block comment right above it (modulo the variable declaration that is only used inside the lambda), so I wouldn't call it auto-documenting since you felt like you needed to write documentation for it. And having a named variable of lambda-type or an identically-named variable of boolean type still gives you exactly the same naming scheme.

Given @ntv's and @mehdi_amini's comments, I'm now strongly inclined to retain it.

I read Mehdi's comment differently:

but in general not calling it right after

rather to outline a block of code outside of a loop to make the loop shorter

do not seem to necessarily support your usage here (neither does it contradict).

Readability is a very subjective thing. I was reading your code for review purposes and this construct did make me lose time and expand more energy than for a straight-line code here, so for me personally it decreased readability. Namely because it (a) mutates an implicitly captured variable and (b) requires to unwrap more abstractions mentally. Anyway, I won't block the commit just because of a stylistic discussion.

I can suggest an alternative that would address part of my readability concerns:

int64_t constEltSizeBytes = [elementType]() {
  if (auto vectorType = elementType.template dyn_cast<VectorType>())
    return vectorType.getNumElements() *
        llvm::divideCeil(vectorType.getElementTypeBitWidth(), 8);
  else
    return llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
}();
bool isMallocAlignmentSufficient = constEltSizeBytes > kMallocAlignment;

This removes implicit by-reference capture, makes it clear that you do not intend for the lambda to be reused (named lambda would be also okay since it doesn't store references anyway, but there's no point), and this way of using lambdas is actually considered a C++11 idiom for comlex constant initialization (https://groups.google.com/a/isocpp.org/g/std-discussion/c/FBjcR4WJlkU/m/nQnsSOziq04J) so one can claim it's "common enough".

1582

How about just keeping it simple with -use-aligned-alloc and not changing previous/existing behavior when -use-aligned-alloc is not provided?

Works for me!

[rriddle]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1420

FWIWI, there is already a guide on using lamdas for computing predicates:

https://llvm.org/docs/CodingStandards.html#turn-predicate-loops-into-predicate-functions

[mehdi_amini]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1420

While I haven't seen this pattern used as it is here widely in MLIR (and I like consistency), I just didn't feel any reason for me to really oppose it, but I wasn't really supporting it as-is either. The fact that @ftynse opposed to this on principle of readability seems to hint that it would be better to not do it.

I would like have written it in a form around this (with a comment possibly):

int64_t constEltSizeBytes = llvm::divideCeil(vectorType.getElementTypeBitWidth(), 8);
if (auto vectorType = elementType.template dyn_cast<VectorType>())
  constEltSizeBytes = constEltSizeBytes * vectorType.getNumElements();
bool isMallocAlignmentSufficient = constEltSizeBytes > kMallocAlignment;

Everything is named as well and you avoid the readability overhead of the lambda.

1420

@rriddle: the link is about free functions I believe, which I see differently than lambda: the rational explained there is close to what I mention before of moving the predicate code "out of the way" of the main logic (and reducing the indentation, etc), which you don't get with the immediately local lambda.

[bondhugula]

@ftynse Everything's addressed now but you may want to take another look because I had completely missed the fact earlier that aligned_alloc only supports a size that is a multiple of alignment. So I have additional handling (and test cases for that) now to bump the allocation size to the next multiple of alignment. Also, if the elt size is not a power of two (and an alignment attribute doesn't exist), the next power of two is used for alignment (although this isn't an interesting case for aligned allocation, didn't want to punt to malloc for this to keep things simple/clear). So, all heap allocations use aligned_alloc whenever -use-aligned-alloc is set.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1408

All of this now dropped. -use-aligned-alloc is now independent of what malloc can do. The malloc path remains untouched.

1420

Thank you all for commenting. This discussion is now moot since the aligned_alloc generation is now untangled from being contingent on this conditional / what malloc supports, and so this whole lambda and conditionals are gone.

1582

Done - aligned_alloc is now used only with -use-aligned-alloc, and for all heap allocations whenever that cmd line flag exists.

[bondhugula]

Address review comments. Also, handle when the size may not be a multiple of alignment.

[ftynse]

Everything's addressed now but you may want to take another look because I had completely missed the fact earlier that aligned_alloc only supports a size that is a multiple of alignment.

I was going to ask what happened in case of memref<?xvector<3xf32>> but got distracted by the lambda discussion :)

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1584

Typo: Aigned

[bondhugula]

Add another test case.

[bondhugula]

In D77528#1968908, @ftynse wrote:

Everything's addressed now but you may want to take another look because I had completely missed the fact earlier that aligned_alloc only supports a size that is a multiple of alignment.

I was going to ask what happened in case of memref<?xvector<3xf32>> but got distracted by the lambda discussion :)

The allocation size will be bumped to a multiple of the alignment. With memref<?xvector<3xf32>, if there is no alignment attribute, the next power of two larger than 12, i.e., 16, will be the alignment, and the allocation size will be bumped to a multiple of 16 ('cumulativeSize' will bumped up).

[bondhugula]

Address review comments.

[ftynse]

This seems to have broken the build with gcc-5, https://buildkite.com/mlir/mlir-core/builds/4023#0cd4c474-dbf0-4b26-a3f5-d2a63b983ac9.

[bondhugula]

In D77528#1969067, @ftynse wrote:

This seems to have broken the build with gcc-5, https://buildkite.com/mlir/mlir-core/builds/4023#0cd4c474-dbf0-4b26-a3f5-d2a63b983ac9.

I don't have GCC-5 at hand to be sure of the fix. Any insights?

[ftynse]

In D77528#1969077, @bondhugula wrote:

I don't have GCC-5 at hand to be sure of the fix. Any insights?

Neither do I... I would expect using AllocaOpLowering = AllocLikeOpLowering<AllocaOp>; instead of creating a derived class could get rid of the problem with constructors, and this is the right thing to have anyway.

[bondhugula]

In D77528#1969104, @ftynse wrote:

In D77528#1969077, @bondhugula wrote:

I don't have GCC-5 at hand to be sure of the fix. Any insights?

Neither do I... I would expect using AllocaOpLowering = AllocLikeOpLowering<AllocaOp>; instead of creating a derived class could get rid of the problem with constructors, and this is the right thing to have anyway.

Thanks - I just went ahead and committed this: D77719. Where can I get the buildkite URL for a commit?

[ftynse]

In D77528#1969205, @bondhugula wrote:

Thanks - I just went ahead and committed this: D77719. Where can I get the buildkite URL for a commit?

Thanks! The build is performed hourly, using all the commits landed in the last hour, so we'll have to wait a bit. The overall status is here https://buildkite.com/mlir/mlir-core

[bondhugula]

In D77528#1969290, @ftynse wrote:

In D77528#1969205, @bondhugula wrote:

Thanks - I just went ahead and committed this: D77719. Where can I get the buildkite URL for a commit?

Thanks! The build is performed hourly, using all the commits landed in the last hour, so we'll have to wait a bit. The overall status is here https://buildkite.com/mlir/mlir-core

That didn't work. I made a couple more cleanup changes in D77726 and that has fixed it.

[mehdi_amini]

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

186

Can you split this in a new test file?

We will process the entire file twice while the tests between the two invocations are actually entirely disjoint.

[bondhugula]

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

186

Sure - this make sense. Given a similar approach and a tendency to put everything in one file at other places, we missed that. The only benefit of having it here is that the generated ocde follows the pattern of the test case right above and so it was easy to add it here and update both together when needed. But thinking about this: what about the additional proofing this provides given that it is also running (without crashing, etc.) on the other cases even if we don't have a FileCheck for the rest? I think there is some benefit there without having to copy over that stuff to another file? How do you weigh these?

mlir/include/mlir/Conversion/Passes.td

let description = [{
     Functions converted to LLVM IR. Function arguments types are converted
     one-to-one. Function results are converted one-to-one and, in case more than
     1 value is returned, packed into an LLVM IR struct type. Function calls and
     returns are updated accordingly. Block argument types are updated to use
     LLVM IR types.
   }];
   let constructor = "mlir::createLowerToLLVMPass()";
   let options = [
-    Option<"useAlloca", "use-alloca", "bool", /*default=*/"false",
-           "Use `alloca` instead of `call @malloc` for converting std.alloc">,
+    Option<"alignedAlloc", "aligned-alloc", "bool", /*default=*/"false",
+           "Use aligned_alloc in place of malloc whenever alignment is set on "
+           "std.alloc">,
     Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
            /*default=*/"false",
            "Replace FuncOp's MemRef arguments with bare pointers to the MemRef "
            "element types">,
     Option<"emitCWrappers", "emit-c-wrappers", "bool", /*default=*/"false",
            "Emit wrappers for C-compatible pointer-to-struct memref "
            "descriptors">,
     Option<"indexBitwidth", "index-bitwidth", "unsigned",

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

 //===- ConvertStandardToLLVMPass.h - Pass entrypoint ------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_
 #define MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_
 
 #include <memory>
 
 namespace mlir {
 class LLVMTypeConverter;
 class ModuleOp;
 template <typename T> class OpPassBase;
 class OwningRewritePatternList;
 
 /// Collect a set of patterns to convert memory-related operations from the
 /// Standard dialect to the LLVM dialect, excluding non-memory-related
 /// operations and FuncOp.
-void populateStdToLLVMMemoryConversionPatters(
-    LLVMTypeConverter &converter, OwningRewritePatternList &patterns);
+void populateStdToLLVMMemoryConversionPatterns(
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
+    bool alignedAlloc);
 
 /// Collect a set of patterns to convert from the Standard dialect to the LLVM
 /// dialect, excluding the memory-related operations.
 void populateStdToLLVMNonMemoryConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns);
 
 /// Collect the default pattern to convert a FuncOp to the LLVM dialect. If
 /// `emitCWrappers` is set, the pattern will also produce functions
 /// that pass memref descriptors by pointer-to-structure in addition to the
 /// default unpacked form.
 void populateStdToLLVMDefaultFuncOpConversionPattern(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     bool emitCWrappers = false);
 
 /// Collect a set of default patterns to convert from the Standard dialect to
-/// LLVM.
+/// LLVM. If `alignedAlloc` is set, aligned_alloc is used in place of malloc
+/// whenever an alignment is specified.
 void populateStdToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                          OwningRewritePatternList &patterns,
+                                         bool alignedAlloc = false,
                                          bool emitCWrappers = false);
 
 /// Collect a set of patterns to convert from the Standard dialect to
 /// LLVM using the bare pointer calling convention for MemRef function
-/// arguments.
+/// arguments. If `alignedAlloc` is set, aligned_alloc is used in place of
+/// malloc whenever an alignment is specified.
 void populateStdToLLVMBarePtrConversionPatterns(
-    LLVMTypeConverter &converter, OwningRewritePatternList &patterns);
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
+    bool alignedAlloc);
 
 /// Value to pass as bitwidth for the index type when the converter is expected
 /// to derive the bitwidth from the LLVM data layout.
 static constexpr unsigned kDeriveIndexBitwidthFromDataLayout = 0;
 
 /// Creates a pass to convert the Standard dialect into the LLVMIR dialect.
 /// By default stdlib malloc/free are used for allocating MemRef payloads.
+/// Specifying `aligned-alloc-true` uses aligned_alloc instead of malloc to
+/// perform memref allocations whenever an alignment is
+/// specified.
 std::unique_ptr<OpPassBase<ModuleOp>> createLowerToLLVMPass(
-    bool useBarePtrCallConv = false, bool emitCWrappers = false,
+    bool alignedAlloc = false, bool useBarePtrCallConv = false,
+    bool emitCWrappers = false,
     unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout);
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 //===- StandardToLLVM.cpp - Standard to LLVM dialect conversion -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 // Extract an LLVM IR type from the LLVM IR dialect type.
 static LLVM::LLVMType unwrap(Type type) {
   if (!type)
     return nullptr;
   auto *mlirContext = type.getContext();
   auto wrappedLLVMType = type.dyn_cast<LLVM::LLVMType>();
   if (!wrappedLLVMType)
     emitError(UnknownLoc::get(mlirContext),
               "conversion resulted in a non-LLVM type");
   return wrappedLLVMType;
 }
 
 /// Initialize customization to default callbacks.
 LLVMTypeConverterCustomization::LLVMTypeConverterCustomization()
     : funcArgConverter(structFuncArgTypeConverter),
       indexBitwidth(kDeriveIndexBitwidthFromDataLayout) {}
 

 struct AllocLikeOpLowering : public ConvertOpToLLVMPattern<AllocLikeOp> {
   using ConvertOpToLLVMPattern<AllocLikeOp>::ConvertOpToLLVMPattern;
   using Base = AllocLikeOpLowering<AllocLikeOp>;
   using ConvertOpToLLVMPattern<AllocLikeOp>::createIndexConstant;
   using ConvertOpToLLVMPattern<AllocLikeOp>::getIndexType;
   using ConvertOpToLLVMPattern<AllocLikeOp>::typeConverter;
   using ConvertOpToLLVMPattern<AllocLikeOp>::getVoidPtrType;
 
-  explicit AllocLikeOpLowering(LLVMTypeConverter &converter)
-      : ConvertOpToLLVMPattern<AllocLikeOp>(converter) {}
+  explicit AllocLikeOpLowering(LLVMTypeConverter &converter,
+                               bool alignedAlloc = false)
+      : ConvertOpToLLVMPattern<AllocLikeOp>(converter),
+        alignedAlloc(alignedAlloc) {}
 
   LogicalResult match(Operation *op) const override {
     MemRefType memRefType = cast<AllocLikeOp>(op).getType();
     if (isSupportedMemRefType(memRefType))
       return success();
 
     int64_t offset;
     SmallVector<int64_t, 4> strides;
struct AllocLikeOpLowering : public ConvertOpToLLVMPattern<AllocLikeOp> {
   /// Returns the type of a pointer to an element of the memref.
   Type getElementPtrType(MemRefType memRefType) const {
     auto elementType = memRefType.getElementType();
     auto structElementType = typeConverter.convertType(elementType);
     return structElementType.template cast<LLVM::LLVMType>().getPointerTo(
         memRefType.getMemorySpace());
   }
 
+  /// Returns the alignment to be used for the allocation itself.
+  /// `useAlignedAlloc` is set to true if aligned_alloc should be used instead
+  /// of malloc.
+  Optional<int64_t> getAllocationAlignment(Operation *op,
+                                           bool &useAlignedAlloc) const {
+    auto elementType = cast<AllocLikeOp>(op).getType().getElementType();
+
+    useAlignedAlloc = false;
+
+    // For stack allocations.
+    if (auto allocaOp = dyn_cast<AllocaOp>(op)) {
+      auto alignment = cast<AllocaOp>(op).alignment();
+      return alignment ? Optional<int64_t>(alignment.getValue().getSExtValue())
+                       : None;
+    }
+
+    // For heap allocations.
+    AllocOp allocOp = cast<AllocOp>(op);
+    if (alignedAlloc && allocOp.alignment()) {
+      useAlignedAlloc = true;
+      return allocOp.alignment().getValue().getSExtValue();
+    }
+
+    // The alloc op might have an alignment attribute, but there is no alignment
+    // to be used for the malloc call itself.
+    if (allocOp.alignment())
+      return None;
+
+    // Whenever we don't have alignment set, we will still use aligned_alloc
+    // if the alignment needed is more than what malloc can provide.
+    uint64_t constEltSizeBytes = 0;
+    auto isMallocAlignmentSufficient = [&]() {
+      if (auto vectorType = elementType.template dyn_cast<VectorType>())
+        constEltSizeBytes =
+            vectorType.getNumElements() *
+            llvm::divideCeil(vectorType.getElementTypeBitWidth(), 8);
+      else
+        constEltSizeBytes =
+            llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
+
+      // Use aligned_alloc if elt_size > malloc's alignment.
+      return (constEltSizeBytes > kMallocAlignment);
+    };
+    useAlignedAlloc = isMallocAlignmentSufficient();
+    if (useAlignedAlloc)
+      return std::max(constEltSizeBytes,
+                      static_cast<uint64_t>(kMallocAlignment));
+    return None;
+  }
+
   /// Allocates the underlying buffer using the right call. `allocatedBytePtr`
   /// is set to null for stack allocations. `accessAlignment` is set if
   /// alignment is neeeded post allocation (for eg. in conjunction with malloc).
-  /// TODO(bondhugula): next revision will support std lib func aligned_alloc.
   Value allocateBuffer(Location loc, Value cumulativeSize, Operation *op,
                        MemRefType memRefType, Value one, Value &accessAlignment,
                        Value &allocatedBytePtr,
                        ConversionPatternRewriter &rewriter) const {
     auto elementPtrType = getElementPtrType(memRefType);
 
     // Whether to use std lib function aligned_alloc that supports alignment.
-    Optional<APInt> allocationAlignment = cast<AllocLikeOp>(op).alignment();
+    bool useAlignedAlloc;
+    Optional<int64_t> allocationAlignment =
+        getAllocationAlignment(op, useAlignedAlloc);
 
     // With alloca, one gets a pointer to the element type right away.
     bool onStack = isa<AllocaOp>(op);
     if (onStack) {
       allocatedBytePtr = nullptr;
       accessAlignment = nullptr;
       return rewriter.create<LLVM::AllocaOp>(
           loc, elementPtrType, cumulativeSize,
-          allocationAlignment ? allocationAlignment.getValue().getSExtValue()
-                              : 0);
+          allocationAlignment ? allocationAlignment.getValue() : 0);
     }
 
-    // Use malloc. Insert the malloc declaration if it is not already present.
-    auto allocFuncName = "malloc";
+    // Use malloc if useAlignedAlloc is false, otherwise use aligned_alloc;
+
+    // Insert the malloc/aligned_alloc declaration if it is not already present.
+    auto allocFuncName = useAlignedAlloc ? "aligned_alloc" : "malloc";
     AllocOp allocOp = cast<AllocOp>(op);
     auto module = allocOp.getParentOfType<ModuleOp>();
     auto allocFunc = module.lookupSymbol<LLVM::LLVMFuncOp>(allocFuncName);
     if (!allocFunc) {
       OpBuilder moduleBuilder(op->getParentOfType<ModuleOp>().getBodyRegion());
       SmallVector<LLVM::LLVMType, 2> callArgTypes = {getIndexType()};
+      // aligned_alloc(size_t alignment, size_t size)
+      if (useAlignedAlloc)
+        callArgTypes.push_back(getIndexType());
       allocFunc = moduleBuilder.create<LLVM::LLVMFuncOp>(
           rewriter.getUnknownLoc(), allocFuncName,
           LLVM::LLVMType::getFunctionTy(getVoidPtrType(), callArgTypes,
                                         /*isVarArg=*/false));
     }
 
     // Allocate the underlying buffer and store a pointer to it in the MemRef
     // descriptor.
     SmallVector<Value, 2> callArgs;
+    if (useAlignedAlloc) {
+      assert(allocationAlignment && "allocation alignment should be present");
+      // Use aligned_alloc: we don't need to set 'align'.
+      auto alignedAllocAlignmentValue = rewriter.create<LLVM::ConstantOp>(
+          loc, typeConverter.convertType(rewriter.getIntegerType(64)),
+          rewriter.getI64IntegerAttr(allocationAlignment.getValue()));
+      callArgs = {alignedAllocAlignmentValue, cumulativeSize};
+    } else {
       // Adjust the allocation size to consider alignment.
       if (allocOp.alignment()) {
         accessAlignment = createIndexConstant(
             rewriter, loc, allocOp.alignment().getValue().getSExtValue());
         cumulativeSize = rewriter.create<LLVM::SubOp>(
             loc,
             rewriter.create<LLVM::AddOp>(loc, cumulativeSize, accessAlignment),
             one);
       }
       callArgs.push_back(cumulativeSize);
+    }
     auto allocFuncSymbol = rewriter.getSymbolRefAttr(allocFunc);
     allocatedBytePtr = rewriter
                            .create<LLVM::CallOp>(loc, getVoidPtrType(),
                                                  allocFuncSymbol, callArgs)
                            .getResult(0);
     // For heap allocations, the allocated pointer is a cast of the byte pointer
     // to the type pointer.
     return rewriter.create<LLVM::BitcastOp>(loc, elementPtrType,
void rewrite(Operation *op, ArrayRef<Value> operands,
     // Create the MemRef descriptor.
     auto memRefDescriptor = createMemRefDescriptor(
         loc, rewriter, memRefType, allocatedTypePtr, allocatedBytePtr,
         accessAlignment, offset, strides, sizes);
 
     // Return the final value of the descriptor.
     rewriter.replaceOp(op, {memRefDescriptor});
   }
+
+protected:
+  // This is the alignment malloc typically provides.
+  constexpr static unsigned kMallocAlignment = 16;
+  bool alignedAlloc;
 };
 
 struct AllocOpLowering : public AllocLikeOpLowering<AllocOp> {
-  using Base::Base;
+  explicit AllocOpLowering(LLVMTypeConverter &converter,
+                           bool alignedAlloc = false)
+      : AllocLikeOpLowering<AllocOp>(converter, alignedAlloc) {}
 };
+
 struct AllocaOpLowering : public AllocLikeOpLowering<AllocaOp> {
   using Base::Base;
 };
 
 // A CallOp automatically promotes MemRefType to a sequence of alloca/store and
 // passes the pointer to the MemRef across function boundaries.
 template <typename CallOpType>
 struct CallOpInterfaceLowering : public ConvertOpToLLVMPattern<CallOpType> {
patterns.insert<
       UnsignedDivIOpLowering,
       UnsignedRemIOpLowering,
       UnsignedShiftRightOpLowering,
       XOrOpLowering,
       ZeroExtendIOpLowering>(converter);
   // clang-format on
 }
 
-void mlir::populateStdToLLVMMemoryConversionPatters(
-    LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
+void mlir::populateStdToLLVMMemoryConversionPatterns(
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
+    bool alignedAlloc) {
   // clang-format off
   patterns.insert<
       AssumeAlignmentOpLowering,
+      DeallocOpLowering,
       DimOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       StoreOpLowering,
       SubViewOpLowering,
       ViewOpLowering>(converter);
   patterns.insert<
-      AllocOpLowering,
-      DeallocOpLowering>(converter);
+      AllocOpLowering
+      >(converter, alignedAlloc);
   // clang-format on
 }
 
 void mlir::populateStdToLLVMDefaultFuncOpConversionPattern(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     bool emitCWrappers) {
   patterns.insert<FuncOpConversion>(converter, emitCWrappers);
 }
 
 void mlir::populateStdToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
-    bool emitCWrappers) {
+    bool alignedAlloc, bool emitCWrappers) {
   populateStdToLLVMDefaultFuncOpConversionPattern(converter, patterns,
                                                   emitCWrappers);
   populateStdToLLVMNonMemoryConversionPatterns(converter, patterns);
-  populateStdToLLVMMemoryConversionPatters(converter, patterns);
+  populateStdToLLVMMemoryConversionPatterns(converter, patterns, alignedAlloc);
 }
 
 static void populateStdToLLVMBarePtrFuncOpConversionPattern(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   patterns.insert<BarePtrFuncOpConversion>(converter);
 }
 
 void mlir::populateStdToLLVMBarePtrConversionPatterns(
-    LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
+    bool alignedAlloc) {
   populateStdToLLVMBarePtrFuncOpConversionPattern(converter, patterns);
   populateStdToLLVMNonMemoryConversionPatterns(converter, patterns);
-  populateStdToLLVMMemoryConversionPatters(converter, patterns);
+  populateStdToLLVMMemoryConversionPatterns(converter, patterns, alignedAlloc);
 }
 
 // Create an LLVM IR structure type if there is more than one result.
 Type LLVMTypeConverter::packFunctionResults(ArrayRef<Type> types) {
   assert(!types.empty() && "expected non-empty list of type");
 
   if (types.size() == 1)
     return convertType(types.front());

 namespace {
 /// A pass converting MLIR operations into the LLVM IR dialect.
 struct LLVMLoweringPass : public ModulePass<LLVMLoweringPass> {
 /// Include the generated pass utilities.
 #define GEN_PASS_ConvertStandardToLLVM
 #include "mlir/Conversion/Passes.h.inc"
 
   /// Creates an LLVM lowering pass.
-  LLVMLoweringPass(bool useBarePtrCallConv, bool emitCWrappers,
-                   unsigned indexBitwidth) {
+  LLVMLoweringPass(bool alignedAlloc, bool useBarePtrCallConv,
+                   bool emitCWrappers, unsigned indexBitwidth) {
+    this->alignedAlloc = alignedAlloc;
     this->useBarePtrCallConv = useBarePtrCallConv;
     this->emitCWrappers = emitCWrappers;
     this->indexBitwidth = indexBitwidth;
   }
   explicit LLVMLoweringPass() {}
   LLVMLoweringPass(const LLVMLoweringPass &pass) {}
 
   /// Run the dialect converter on the module.
void runOnModule() override {
     LLVMTypeConverterCustomization customs;
     customs.funcArgConverter = useBarePtrCallConv ? barePtrFuncArgTypeConverter
                                                   : structFuncArgTypeConverter;
     customs.indexBitwidth = indexBitwidth;
     LLVMTypeConverter typeConverter(&getContext(), customs);
 
     OwningRewritePatternList patterns;
     if (useBarePtrCallConv)
-      populateStdToLLVMBarePtrConversionPatterns(typeConverter, patterns);
+      populateStdToLLVMBarePtrConversionPatterns(typeConverter, patterns,
+                                                 alignedAlloc);
     else
-      populateStdToLLVMConversionPatterns(typeConverter, patterns,
+      populateStdToLLVMConversionPatterns(typeConverter, patterns, alignedAlloc,
                                           emitCWrappers);
 
     LLVMConversionTarget target(getContext());
     if (failed(applyPartialConversion(m, target, patterns, &typeConverter)))
       signalPassFailure();
   }
 };
 } // end namespace
 
 mlir::LLVMConversionTarget::LLVMConversionTarget(MLIRContext &ctx)
     : ConversionTarget(ctx) {
   this->addLegalDialect<LLVM::LLVMDialect>();
   this->addIllegalOp<LLVM::DialectCastOp>();
   this->addIllegalOp<TanhOp>();
 }
 
 std::unique_ptr<OpPassBase<ModuleOp>>
-mlir::createLowerToLLVMPass(bool useBarePtrCallConv, bool emitCWrappers,
-                            unsigned indexBitwidth) {
-  return std::make_unique<LLVMLoweringPass>(useBarePtrCallConv, emitCWrappers,
-                                            indexBitwidth);
+mlir::createLowerToLLVMPass(bool alignedAlloc, bool useBarePtrCallConv,
+                            bool emitCWrappers, unsigned indexBitwidth) {
+  return std::make_unique<LLVMLoweringPass>(alignedAlloc, useBarePtrCallConv,
+                                            emitCWrappers, indexBitwidth);
 }

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

 // RUN: mlir-opt -convert-std-to-llvm %s | FileCheck %s
+// RUN: mlir-opt -convert-std-to-llvm='aligned-alloc=1' %s | FileCheck %s --check-prefix=ALIGNED-ALLOC
 
 // CHECK-LABEL: func @check_strided_memref_arguments(
 // CHECK-COUNT-2: !llvm<"float*">
 // CHECK-COUNT-5: !llvm.i64
 // CHECK-COUNT-2: !llvm<"float*">
 // CHECK-COUNT-5: !llvm.i64
 // CHECK-COUNT-2: !llvm<"float*">
 // CHECK-COUNT-5: !llvm.i64

 func @dynamic_dealloc(%arg0: memref<?x?xf32>) {
 //      CHECK:  %[[ptr:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
 // CHECK-NEXT:  %[[ptri8:.*]] = llvm.bitcast %[[ptr]] : !llvm<"float*"> to !llvm<"i8*">
 // CHECK-NEXT:  llvm.call @free(%[[ptri8]]) : (!llvm<"i8*">) -> ()
   dealloc %arg0 : memref<?x?xf32>
   return
 }
 
+// CHECK-LABEL: func @stdlib_aligned_alloc() -> !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }"> {
+// ALIGNED-ALLOC-LABEL: func @stdlib_aligned_alloc() -> !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }"> {
+func @stdlib_aligned_alloc() -> memref<32x18xf32> {
+// ALIGNED-ALLOC-NEXT:  %[[sz1:.*]] = llvm.mlir.constant(32 : index) : !llvm.i64
+// ALIGNED-ALLOC-NEXT:  %[[sz2:.*]] = llvm.mlir.constant(18 : index) : !llvm.i64
+// ALIGNED-ALLOC-NEXT:  %[[num_elems:.*]] = llvm.mul %0, %1 : !llvm.i64
+// ALIGNED-ALLOC-NEXT:  %[[null:.*]] = llvm.mlir.null : !llvm<"float*">
+// ALIGNED-ALLOC-NEXT:  %[[one:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+// ALIGNED-ALLOC-NEXT:  %[[gep:.*]] = llvm.getelementptr %[[null]][%[[one]]] : (!llvm<"float*">, !llvm.i64) -> !llvm<"float*">
+// ALIGNED-ALLOC-NEXT:  %[[sizeof:.*]] = llvm.ptrtoint %[[gep]] : !llvm<"float*"> to !llvm.i64
+// ALIGNED-ALLOC-NEXT:  %[[bytes:.*]] = llvm.mul %[[num_elems]], %[[sizeof]] : !llvm.i64
+// ALIGNED-ALLOC-NEXT:  %[[alignment:.*]] = llvm.mlir.constant(32 : i64) : !llvm.i64
+// ALIGNED-ALLOC-NEXT:  %[[allocated:.*]] = llvm.call @aligned_alloc(%[[alignment]], %[[bytes]]) : (!llvm.i64, !llvm.i64) -> !llvm<"i8*">
+// ALIGNED-ALLOC-NEXT:  llvm.bitcast %[[allocated]] : !llvm<"i8*"> to !llvm<"float*">
+  %0 = alloc() {alignment = 32} : memref<32x18xf32>
+  // Do another alloc just to test that we have a unique declaration for
+  // aligned_alloc.
+  // ALIGNED-ALLOC:  llvm.call @aligned_alloc
+  %1 = alloc() {alignment = 64} : memref<4096xf32>
+
+  // Memrefs with elements types larger than 16 bytes are allocated via
+  // aligned_alloc even without an alignment attribute.
+  // ALIGNED-ALLOC:  llvm.call @aligned_alloc
+  %2 = alloc() : memref<4096xvector<8xf32>>
+  // While those with elt type <= 16 bytes still use malloc unless the alignment
+  // attribute is specified.
+  // ALIGNED-ALLOC:  llvm.call @malloc
+  %3 = alloc() : memref<4096xvector<2xf32>>
+  // ALIGNED-ALLOC:  llvm.call @aligned_alloc
+  %4 = alloc() {alignment = 16} : memref<4096xvector<2xf32>>
+  return %0 : memref<32x18xf32>
+}
+
 // CHECK-LABEL: func @mixed_load(
 // CHECK-COUNT-2: !llvm<"float*">,
 // CHECK-COUNT-5: {{%[a-zA-Z0-9]*}}: !llvm.i64
 // CHECK:         %[[I:.*]]: !llvm.i64,
 // CHECK:         %[[J:.*]]: !llvm.i64)
 func @mixed_load(%mixed : memref<42x?xf32>, %i : index, %j : index) {
 //       CHECK:  %[[ptr:.*]] = llvm.extractvalue %[[ld:.*]][1] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
 //  CHECK-NEXT:  %[[off:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
 //  CHECK-NEXT:  %[[st0:.*]] = llvm.extractvalue %[[ld]][4, 0] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
 //  CHECK-NEXT:  %[[offI:.*]] = llvm.mul %[[I]], %[[st0]] : !llvm.i64
 //  CHECK-NEXT:  %[[off0:.*]] = llvm.add %[[off]], %[[offI]] : !llvm.i64
 //  CHECK-NEXT:  %[[st1:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
 //  CHECK-NEXT:  %[[offJ:.*]] = llvm.mul %[[J]], %[[st1]] : !llvm.i64
 //  CHECK-NEXT:  %[[off1:.*]] = llvm.add %[[off0]], %[[offJ]] : !llvm.i64
 //  CHECK-NEXT:  %[[addr:.*]] = llvm.getelementptr %[[ptr]][%[[off1]]] : (!llvm<"float*">, !llvm.i64) -> !llvm<"float*">
 //  CHECK-NEXT:  llvm.load %[[addr]] : !llvm<"float*">
   %0 = load %mixed[%i, %j] : memref<42x?xf32>
   return
 }