[mlir] Add use-vector-alignment flag to ConvertVectorToLLVMPass #137389

electriclilies · 2025-04-25T20:03:31Z

In ConvertVectorToLLVM, the only option for setting alignment of vector.gather, vector.scatter, and the vector.load/store ops was to extract it from the datatype of the memref type. However, this is insufficient for hardware backends requiring alignment of vector types. This PR introduces the use-vector-alignment option to the ConvertVectorToLLVMPass, which makes the pass use the alignment of the vector type of these operations instead of the alignment of the memref type.

llvmbot · 2025-04-25T20:04:06Z

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Lily Orth-Smith (electriclilies)

Changes

Previously, the alignment of a MemRefType was set to the alignment of its element type. This caused us to have to set the preferred alignment of scalar types to the desired preferred alignment for vector types. For example, we had to set the preferred alignment of i1s to 512 in our DataLayoutDescription, which is not ideal.

This change constructs a vector type which has the same datatype and size as the MemRefType, and then extracts the preferred alignment from that vector type. I did try to use the MemRefType directly, but the output alignments were not correct.

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

1 Files Affected:

(modified) mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp (+5-2)

diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index 076e5512f375b..480819bd07680 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -70,8 +70,11 @@ static Value extractOne(ConversionPatternRewriter &rewriter,
 // Helper that returns data layout alignment of a memref.
 LogicalResult getMemRefAlignment(const LLVMTypeConverter &typeConverter,
                                  MemRefType memrefType, unsigned &align) {
-  Type elementTy = typeConverter.convertType(memrefType.getElementType());
-  if (!elementTy)
+  // Align MemRefTypes to the alignment of a VectorType with the same
+  // size and dtype.
+  Type convertedVecType = typeConverter.convertType(
+      VectorType::get(memrefType.getShape(), memrefType.getElementType()));
+  if (!convertedVecType)
     return failure();
 
   // TODO: this should use the MLIR data layout when it becomes available and

github-actions · 2025-04-25T22:11:17Z

✅ With the latest revision this PR passed the C/C++ code formatter.

banach-space · 2025-04-28T15:19:51Z

I can see how this can solve your problem for i1, but what should happen for i8 or i32? Importantly, this only solves your problem for statically-shaped MemRef(s), right?

Why not align vectors of sub-bytes on byte boundary?

Btw, we need a way to test this change. MLIR examples would also be helpful (e.g. in the summary).

electriclilies · 2025-04-28T16:17:15Z

At the point in the stack we use these patterns, we've gotten rid of all non-statically shaped MemRefs. But I can see how this would be a problem for the generic use case.

I can't reveal too many details about our hardware, but most SIMD vectors in our hardware backend have a specific preferred alignment that we need to specify. For example, we had to set the preferred alignment of f16, f32, and n32 to the preferred alignment for their corresponding vector dtypes. So a portion of our data layout description looks like this, because the preferred alignments for vector types aren't being used

f16:{vector_alignment}-f32:{vector_alignment}-n32:{vector_alignment}.

This isn't ideal because these alignments are also being used for scalar alignment. What we'd like is for our data layout description to look more like this:

f16:{f16_alignment}-f32:{f32_alignment}-n32:{n32_alignment}-v512:{vector_alignment}-v1024:{vector_alignment}-v2048:{vector_alignment}-v4096:{vector_alignment}-v8192

electriclilies · 2025-04-28T16:22:35Z

I agree that we need a way to test this upstream. That would probably require adding a test that has a different or fake data layout description, and I'm not sure how to do that.

banach-space · 2025-04-29T19:53:47Z

Hey @electriclilies — this isn’t my area of expertise, but since no one else has chimed in yet, I’ll try to help (at least by pointing you toward the right people).

Right now, this feels a bit ad-hoc and assumes that everyone wants their MemRef(s) aligned to a vector boundary rather than a scalar one. That’s not ideal and should generally be avoided — but having a flag to control this seems fine. To me, your requirements don’t seem that exotic and should be relatively straightforward to support. Would adding a flag be sufficient for your use case?

What we'd like is for our data layout description to look more like this:

f16:{f16_alignment}-f32:{f32_alignment}-n32:{n32_alignment}-v512:{vector_alignment}-v1024:{vector_alignment}-v2048:{vector_alignment}-v4096:{vector_alignment}-v8192

IIUC, this format is supported by LLVM’s data layout:

https://llvm.org/docs/LangRef.html#data-layout

Is there a specific issue you’re running into with this?

That would probably require adding a test that has a different or fake data layout description, and I'm not sure how to do that.

A couple of questions:

How are you currently specifying the data layout?
If we added a flag to control whether alignment should match the element or vector boundary, could we test the behavior by toggling that flag?

As per specifying the data layout in a test, I will need to ask around myself 😅

electriclilies · 2025-04-29T22:05:15Z

@banach_space Thanks for the response! I was thinking yesterday about the code I wrote here and I agree it is a bit ad-hoc and actually incorrect in some cases.

If we consider this example from vector_to_llvm_interface.mlir, my old implementation would try to calculate the alignment based on the total size of memref<200x100xf32>, since the vector.store op contains the type of the memref itself, not type of the memref slice.

func.func @store_0d(%memref : memref<200x100xf32>, %i : index, %j : index) {
  %val = arith.constant dense<11.0> : vector<f32>
  vector.store %val, %memref[%i, %j] : memref<200x100xf32>, vector<f32>
  return
}

I think adding a flag is a solution that would work well for us. In our case, we always have the vector type specified, so I would prefer to use the vector type directly instead of extracting the data type and size from the memref type, converting that into a vector type, and then getting the alignment.

IIUC, this format is supported by LLVM’s data layout:

Yes, it is supported, but the vector alignments aren't actually being used by MLIR during vector lowering. So specifying the vector alignments doesn't actually cause the generated vectors to be aligned.

For example, when we set the alignment of f32 scalars to 64 bits, the alignment in the code generated below is 8, even if we've set the alignment of 4096 vectors (128 * 32) in the data layout string.

llvm.store %41, %43 {alignment 8 = i64} : vector<128xf32>, !llvm.ptr
In terms of how we are specifying the data layout -- we have a tablegen file representing our target which contains the data layout string. I think we probably set the data layout string in the LLVMContext somewhere, but I'm not very familiar with that code path. I will look into it and get back to you.

I've updated the code to use a flag toggle between using vector alignment and memref alignment-- let me know what you think! I've also added a test for all the ops affected by this change.

electriclilies · 2025-04-29T23:40:21Z

Also, it looks like I can pass a data layout description into mlir-opt as an option! So, problem solved there.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

banach-space

This approach makes sense to me. I've left some comments inline.

Thanks for working on this!

mlir/include/mlir/Conversion/Passes.td

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

mlir/test/Conversion/VectorToLLVM/use-vector-alignment.mlir

banach-space

Just some final (minor) pointers/requests :) Thanks for all the updates so far!

banach-space · 2025-05-01T11:53:18Z

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

@@ -82,6 +97,22 @@ LogicalResult getMemRefAlignment(const LLVMTypeConverter &typeConverter,
  return success();
 }

+LogicalResult getVectorToLLVMAlignment(const LLVMTypeConverter &typeConverter,


Could you document this as well? Some context useVectorAlignment would be helpful (similar to what you did elsewhere).

banach-space · 2025-05-01T11:55:41Z

mlir/test/Conversion/VectorToLLVM/use-vector-alignment.mlir

+// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=0' --split-input-file | FileCheck %s --check-prefix=MEMREF-ALIGN
+// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=1' --split-input-file | FileCheck %s --check-prefix=VEC-ALIGN


Previously this would also specify the actual alignment - what happened here?

I think the option I was passing in was not actually doing anything-- it was setting the data layout as a MLIR attribute at the module level but I guess that wasn't being propagated into the LLVM context? The only test that uses the option just checks that the data layout description attr is present at the module level, and doesn't do anything else with it.

There's a note in the original getMemRefAlignment method saying that we should be getting the data layout description from MLIR but instead it's gotten from the LLVMContext. So I think it's now just using whatever the default is? I'll try and look into it a bit more

Thanks for the explanation and for being so diligent about this. No need to dig deeper, this is sufficient for me.

Yeah - in some future point we should have a MLIR data layout <=> LLVM data layout translation that'll make everything work nicely.

I think you'll get what you want by changing the options(context) to

LowerToLLVMOptions options( ctx, DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));

That'll let you pick things up from an MLIR DLTI, though I don't think that supports vector alignments yet

You can also check for llvm.data_layout attributes (aka LLVM::LLVMDialect::getDataLayoutAttrName()) on the op you're processing's enclosing module and do options.datalayout =

banach-space · 2025-05-01T12:02:44Z

mlir/test/Conversion/VectorToLLVM/use-vector-alignment.mlir

+// VEC-ALIGN-LABEL: func @load
+// VEC-ALIGN: %[[C100:.*]] = llvm.mlir.constant(100 : index) : i64
+// VEC-ALIGN: %[[MUL:.*]] = llvm.mul %{{.*}}, %[[C100]]  : i64
+// VEC-ALIGN: %[[ADD:.*]] = llvm.add %[[MUL]], %{{.*}}  : i64
+// VEC-ALIGN: %[[GEP:.*]] = llvm.getelementptr %{{.*}}[%[[ADD]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+// VEC-ALIGN: llvm.load %[[GEP]] {alignment = 32 : i64} : !llvm.ptr -> vector<8xf32>
+
+// MEMREF-ALIGN-LABEL: func @load
+// MEMREF-ALIGN: %[[C100:.*]] = llvm.mlir.constant(100 : index) : i64
+// MEMREF-ALIGN: %[[MUL:.*]] = llvm.mul %{{.*}}, %[[C100]]  : i64
+// MEMREF-ALIGN: %[[ADD:.*]] = llvm.add %[[MUL]], %{{.*}}  : i64
+// MEMREF-ALIGN: %[[GEP:.*]] = llvm.getelementptr %{{.*}}[%[[ADD]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+// MEMREF-ALIGN: llvm.load %[[GEP]] {alignment = 4 : i64} : !llvm.ptr -> vector<8xf32>


Since we only care about the alignment here (and the rest of the conversion logic is already well tested elsewhere), I think it would be fine to reduce the CHECK lines to something like:

// ALL-LABEL: func @load // VEC-ALIGN: llvm.load %[[GEP]] {alignment = 32 : i64} : !llvm.ptr -> vector<8xf32> // MEMREF-ALIGN: llvm.load %[[GEP]] {alignment = 4 : i64} : !llvm.ptr -> vector<8xf32>

This is mostly in the spirit of "less is more", but it's also consistent with our FileCheck best practices:

"Tests should be minimal, and only check what is absolutely necessary."

Also, note that you can use a dedicated prefix for lines shared between configurations. For example:

llvm-project/mlir/test/Dialect/Vector/linearize.mlir

Line 5 in 082598a

// ALL-LABEL: test_linearize

Thanks for the tips! The file is much smaller now :)

banach-space

LGTM!

This is very nicely polished - well done, and thank you for addressing all my comments 🙏🏻

Just to confirm: do you have commit access? If not, I’m happy to land this for you. Normally, I’d wait a few more days to give other reviewers a chance to chime in, but I’ll be away next week and don’t want to delay you - so I’m happy to merge it tomorrow afternoon (Europe time) if that works.

electriclilies · 2025-05-01T19:16:34Z

Thank you! I do not have commit access-- I did at one point but was inactive for a while and now I don't anymore.
Merging tomorrow afternoon your time works great for me!

krzysz00

A few notes about how to sneak data layout in

(See LowerGPUOpsToROCDL for where I was getting some examples from)

krzysz00 · 2025-05-01T22:53:26Z

mlir/include/mlir/Conversion/Passes.td

+           "bool", /*default=*/"false",
+           "Use the preferred alignment of a vector type in load/store "
+           "operations instead of the alignment of the element type of the "
+           "memref. This flag is intended for use with hardware which requires" 


"or in application contexts where it is known all vector accesses are naturally aligned"?

krzysz00 · 2025-05-01T22:56:16Z

mlir/test/Conversion/VectorToLLVM/use-vector-alignment.mlir

+// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=0' --split-input-file | FileCheck %s --check-prefix=MEMREF-ALIGN
+// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=1' --split-input-file | FileCheck %s --check-prefix=VEC-ALIGN


Yeah - in some future point we should have a MLIR data layout <=> LLVM data layout translation that'll make everything work nicely.

I think you'll get what you want by changing the options(context) to

LowerToLLVMOptions options( ctx, DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));

That'll let you pick things up from an MLIR DLTI, though I don't think that supports vector alignments yet

krzysz00 · 2025-05-01T23:01:40Z

mlir/test/Conversion/VectorToLLVM/use-vector-alignment.mlir

+// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=0' --split-input-file | FileCheck %s --check-prefix=MEMREF-ALIGN
+// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=1' --split-input-file | FileCheck %s --check-prefix=VEC-ALIGN


You can also check for llvm.data_layout attributes (aka LLVM::LLVMDialect::getDataLayoutAttrName()) on the op you're processing's enclosing module and do options.datalayout =

krzysz00 · 2025-05-01T23:08:26Z

Update, on further reading, the MLIR data layout is really not plumbed through - it's only used for index bitwidth. So options.datalayout being set to the nearest llvm.data_layout you find, if any, is the way to go here.

banach-space · 2025-05-02T13:35:35Z

Thanks for all the pointers, these are very helpful! 🙏🏻

Update, on further reading, the MLIR data layout is really not plumbed through - it's only used for index bitwidth. So options.datalayout being set to the nearest llvm.data_layout you find, if any, is the way to go here.

Is it something that you'd like included in this PR? To me, this change is already self-contained, so I am fine with it as is.

electriclilies · 2025-05-02T15:35:43Z

@krzysz00 Thanks for the pointers! I'd prefer to keep this change self contained if that's ok with you. I think your solution would work for the purposes of the test, but I'd like to be able to use the data layout description that's set the LLVMContext, since it's used in the rest of the compiler. To clarify, I've tested with our backend and the vector alignments from the data layout in the LLVMContext are definitely getting correctly propagated.

I also don't like the idea of introducing a new source of truth for the data layout description, especially if we only use it here and not anywhere else in the codebase.

krzysz00 · 2025-05-02T16:04:18Z

Yesh, let's not mess with how the data layout options get set in this PR just for the test.

No need to make the change I suggested, we're good to land

electriclilies · 2025-05-02T16:14:16Z

Great, thanks! I just updated the documentation like you suggested. I don't have merge access so someone else will need to hit the button once it goes green again :)

dcaballe

Great contribution, thanks! I think we should also start thinking about to bring the alignment information to the corresponding vector operations so that eventually we don't have to use this flag.

…#137389) In ConvertVectorToLLVM, the only option for setting alignment of `vector.gather`, `vector.scatter`, and the `vector.load/store` ops was to extract it from the datatype of the memref type. However, this is insufficient for hardware backends requiring alignment of vector types. This PR introduces the `use-vector-alignment` option to the `ConvertVectorToLLVMPass`, which makes the pass use the alignment of the vector type of these operations instead of the alignment of the memref type. --------- Co-authored-by: Lily Orth-Smith <[email protected]>

electriclilies requested review from banach-space, dcaballe and nicolasvasilache as code owners April 25, 2025 20:03

llvmbot added the mlir label Apr 25, 2025

electriclilies requested a review from ftynse April 25, 2025 20:06

electriclilies force-pushed the vector-alignment branch 2 times, most recently from 18dfae0 to c9997ff Compare April 25, 2025 22:07

electriclilies force-pushed the vector-alignment branch from c9997ff to de73656 Compare April 25, 2025 22:15

electriclilies force-pushed the vector-alignment branch 4 times, most recently from 0baf2d2 to 0512c00 Compare April 29, 2025 23:36

llvmbot added mlir:vectorops mlir:vector labels Apr 29, 2025

electriclilies force-pushed the vector-alignment branch 3 times, most recently from d2f5741 to c9cc43c Compare April 30, 2025 00:19

red1bluelost reviewed Apr 30, 2025

View reviewed changes

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp Outdated Show resolved Hide resolved

banach-space reviewed Apr 30, 2025

View reviewed changes

electriclilies force-pushed the vector-alignment branch 2 times, most recently from 077057f to 29598e8 Compare April 30, 2025 20:43

electriclilies changed the title ~~[mlir] Fix MemRefType alignment in ConvertVectorToLLVM~~ [mlir] Add use-vector-alignment flag to ConvertVectorToLLVMPass Apr 30, 2025

electriclilies requested review from banach-space and red1bluelost May 1, 2025 06:44

banach-space reviewed May 1, 2025

View reviewed changes

banach-space approved these changes May 1, 2025

View reviewed changes

Squashed changes

27a25cd

krzysz00 reviewed May 1, 2025

View reviewed changes

Update comment

47a9f63

electriclilies force-pushed the vector-alignment branch from 9ad8538 to 47a9f63 Compare May 2, 2025 16:13

dcaballe approved these changes May 2, 2025

View reviewed changes

sstamenova merged commit 3715de9 into llvm:main May 2, 2025
11 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir] Add use-vector-alignment flag to ConvertVectorToLLVMPass #137389

[mlir] Add use-vector-alignment flag to ConvertVectorToLLVMPass #137389

electriclilies commented Apr 25, 2025 •

edited

Loading

llvmbot commented Apr 25, 2025 •

edited

Loading

github-actions bot commented Apr 25, 2025 •

edited

Loading

banach-space commented Apr 28, 2025

electriclilies commented Apr 28, 2025

electriclilies commented Apr 28, 2025

banach-space commented Apr 29, 2025

electriclilies commented Apr 29, 2025 •

edited

Loading

electriclilies commented Apr 29, 2025

banach-space left a comment

banach-space left a comment

banach-space May 1, 2025

banach-space May 1, 2025

electriclilies May 1, 2025

banach-space May 1, 2025

krzysz00 May 1, 2025

krzysz00 May 1, 2025

banach-space May 1, 2025

electriclilies May 1, 2025

banach-space left a comment

electriclilies commented May 1, 2025

krzysz00 left a comment

krzysz00 May 1, 2025

krzysz00 May 1, 2025

krzysz00 May 1, 2025

krzysz00 commented May 1, 2025

banach-space commented May 2, 2025

electriclilies commented May 2, 2025 •

edited

Loading

krzysz00 commented May 2, 2025

electriclilies commented May 2, 2025

dcaballe left a comment

		// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=0' --split-input-file \| FileCheck %s --check-prefix=MEMREF-ALIGN
		// RUN: mlir-opt %s --convert-vector-to-llvm='use-vector-alignment=1' --split-input-file \| FileCheck %s --check-prefix=VEC-ALIGN

[mlir] Add use-vector-alignment flag to ConvertVectorToLLVMPass #137389

[mlir] Add use-vector-alignment flag to ConvertVectorToLLVMPass #137389

Conversation

electriclilies commented Apr 25, 2025 • edited Loading

llvmbot commented Apr 25, 2025 • edited Loading

github-actions bot commented Apr 25, 2025 • edited Loading

banach-space commented Apr 28, 2025

electriclilies commented Apr 28, 2025

electriclilies commented Apr 28, 2025

banach-space commented Apr 29, 2025

electriclilies commented Apr 29, 2025 • edited Loading

electriclilies commented Apr 29, 2025

banach-space left a comment

Choose a reason for hiding this comment

banach-space left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

banach-space left a comment

Choose a reason for hiding this comment

electriclilies commented May 1, 2025

krzysz00 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

krzysz00 commented May 1, 2025

banach-space commented May 2, 2025

electriclilies commented May 2, 2025 • edited Loading

krzysz00 commented May 2, 2025

electriclilies commented May 2, 2025

dcaballe left a comment

Choose a reason for hiding this comment

electriclilies commented Apr 25, 2025 •

edited

Loading

llvmbot commented Apr 25, 2025 •

edited

Loading

github-actions bot commented Apr 25, 2025 •

edited

Loading

electriclilies commented Apr 29, 2025 •

edited

Loading

electriclilies commented May 2, 2025 •

edited

Loading