[mlir][Vector] Fix vector.transfer alignment calculation

[nicolasvasilache]

Summary

https://reviews.llvm.org/D79246 introduces alignment propagation for vector transfer operations. Unfortunately, the alignment calculation is incorrect and can result in crashes.

This revision fixes the calculation by using the natural alignment of the memref elemental type, instead of the resulting vector type.

If more alignment is desired, it can be done in 2 ways:

1. use a proper vector.type_cast to transform a memref<axbxcxdxf32> into a memref<axbxvector<cxdxf32>> giving a natural alignment of vector<cxdxf32>
2. add an alignment attribute to vector transfer operations and propagate it.

With this change the alignment in the relevant tests goes down from 128 to 4.

[whchung]

Thanks for the patch. On the application I'm working on, the longest vector type would be vector<?x4>, plus I use scalarize-masked-mem-intrin in my pipeline so the issue addressed in this patch was not encountered on my targets before.

[nicolasvasilache]

Add a VectorToSCF pass.

[nicolasvasilache]

Drop spurious reformat.

mlir/include/mlir/Dialect/Vector/VectorOps.h

enum class VectorContractLowering {
   /// Lower to `vector.matrix_multiply`, maps 1-1 to LLVM matrix intrinsics.
   Matmul = 1,
   /// Lower to `vector.outerproduct`.
   OuterProduct = 2,
 };
 /// Structure to control the behavior of vector transform patterns.
 struct VectorTransformsOptions {
   VectorContractLowering vectorContractLowering = VectorContractLowering::FMA;
+  VectorTransformsOptions &
+  setVectorTransformsOptions(VectorContractLowering opt) {
+    vectorContractLowering = opt;
+    return *this;
+  }
 };
 
 /// Collect a set of transformation patterns that are related to contracting
 /// or expanding vector operations:
 ///   ContractionOpLowering,
 ///   ShapeCastOp2DDownCastRewritePattern,
 ///   ShapeCastOp2DUpCastRewritePattern
 ///   BroadcastOpLowering,

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) {
       desc.setStride(rewriter, loc, index, stride);
     }
 
     rewriter.replaceOp(op, {desc});
     return success();
   }
 };
 
-LogicalResult getLLVMTypeAndAlignment(LLVMTypeConverter &typeConverter,
-                                      Type type, LLVM::LLVMType &llvmType,
-                                      unsigned &align) {
-  auto convertedType = typeConverter.convertType(type);
-  if (!convertedType)
+template <typename TransferOp>
+LogicalResult getVectorTransferAlignment(LLVMTypeConverter &typeConverter,
+                                         TransferOp xferOp, unsigned &align) {
+  Type elementTy =
+      typeConverter.convertType(xferOp.getMemRefType().getElementType());
+  if (!elementTy)
     return failure();
 
-  llvmType = convertedType.template cast<LLVM::LLVMType>();
   auto dataLayout = typeConverter.getDialect()->getLLVMModule().getDataLayout();
-  align = dataLayout.getPrefTypeAlignment(llvmType.getUnderlyingType());
+  align = dataLayout.getPrefTypeAlignment(
+      elementTy.cast<LLVM::LLVMType>().getUnderlyingType());
   return success();
 }
 
 LogicalResult
 replaceTransferOpWithLoadOrStore(ConversionPatternRewriter &rewriter,
                                  LLVMTypeConverter &typeConverter, Location loc,
                                  TransferReadOp xferOp,
                                  ArrayRef<Value> operands, Value dataPtr) {
-  LLVM::LLVMType vecTy;
-  unsigned align;
-  if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
-                                     vecTy, align)))
-    return failure();
   rewriter.replaceOpWithNewOp<LLVM::LoadOp>(xferOp, dataPtr);
   return success();
 }
 
 LogicalResult replaceTransferOpWithMasked(ConversionPatternRewriter &rewriter,
                                           LLVMTypeConverter &typeConverter,
                                           Location loc, TransferReadOp xferOp,
                                           ArrayRef<Value> operands,
                                           Value dataPtr, Value mask) {
   auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); };
   VectorType fillType = xferOp.getVectorType();
   Value fill = rewriter.create<SplatOp>(loc, fillType, xferOp.padding());
   fill = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(fillType), fill);
 
-  LLVM::LLVMType vecTy;
+  Type vecTy = typeConverter.convertType(xferOp.getVectorType());
+  if (!vecTy)
+    return failure();
+
   unsigned align;
-  if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
-                                     vecTy, align)))
+  if (failed(getVectorTransferAlignment(typeConverter, xferOp, align)))
     return failure();
 
   rewriter.replaceOpWithNewOp<LLVM::MaskedLoadOp>(
       xferOp, vecTy, dataPtr, mask, ValueRange{fill},
       rewriter.getI32IntegerAttr(align));
   return success();
 }
 
 LogicalResult
 replaceTransferOpWithLoadOrStore(ConversionPatternRewriter &rewriter,
                                  LLVMTypeConverter &typeConverter, Location loc,
                                  TransferWriteOp xferOp,
                                  ArrayRef<Value> operands, Value dataPtr) {
   auto adaptor = TransferWriteOpOperandAdaptor(operands);
-  LLVM::LLVMType vecTy;
-  unsigned align;
-  if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
-                                     vecTy, align)))
-    return failure();
   rewriter.replaceOpWithNewOp<LLVM::StoreOp>(xferOp, adaptor.vector(), dataPtr);
   return success();
 }
 
 LogicalResult replaceTransferOpWithMasked(ConversionPatternRewriter &rewriter,
                                           LLVMTypeConverter &typeConverter,
                                           Location loc, TransferWriteOp xferOp,
                                           ArrayRef<Value> operands,
                                           Value dataPtr, Value mask) {
-  auto adaptor = TransferWriteOpOperandAdaptor(operands);
-  LLVM::LLVMType vecTy;
   unsigned align;
-  if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
-                                     vecTy, align)))
+  if (failed(getVectorTransferAlignment(typeConverter, xferOp, align)))
     return failure();
 
+  auto adaptor = TransferWriteOpOperandAdaptor(operands);
   rewriter.replaceOpWithNewOp<LLVM::MaskedStoreOp>(
       xferOp, adaptor.vector(), dataPtr, mask,
       rewriter.getI32IntegerAttr(align));
   return success();
 }
 
 static TransferReadOpOperandAdaptor
 getTransferOpAdapter(TransferReadOp xferOp, ArrayRef<Value> operands) {

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 //  CHECK-SAME: !llvm<"<17 x i64>">, !llvm<"<17 x i64>">
 //       CHECK: %[[mask:.*]] = llvm.icmp "slt" %[[offsetVec4]], %[[dimVec3]] :
 //  CHECK-SAME: !llvm<"<17 x i64>">
 //
 // 5. Rewrite as a masked read.
 //       CHECK: %[[PASS_THROUGH:.*]] =  llvm.mlir.constant(dense<7.000000e+00> :
 //  CHECK-SAME:  vector<17xf32>) : !llvm<"<17 x float>">
 //       CHECK: %[[loaded:.*]] = llvm.intr.masked.load %[[vecPtr]], %[[mask]],
-//  CHECK-SAME: %[[PASS_THROUGH]] {alignment = 128 : i32} :
+//  CHECK-SAME: %[[PASS_THROUGH]] {alignment = 4 : i32} :
 //  CHECK-SAME: (!llvm<"<17 x float>*">, !llvm<"<17 x i1>">, !llvm<"<17 x float>">) -> !llvm<"<17 x float>">
 
 //
 // 1. Bitcast to vector form.
 //       CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :
 //  CHECK-SAME: (!llvm<"float*">, !llvm.i64) -> !llvm<"float*">
 //       CHECK: %[[vecPtr_b:.*]] = llvm.bitcast %[[gep_b]] :
 //  CHECK-SAME: !llvm<"float*"> to !llvm<"<17 x float>*">

 //       CHECK: llvm.shufflevector {{.*}} [0 : i32, 0 : i32, 0 : i32, 0 : i32,
 //  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32,
 //  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32] :
 //  CHECK-SAME: !llvm<"<17 x i64>">, !llvm<"<17 x i64>">
 //       CHECK: %[[mask_b:.*]] = llvm.icmp "slt" {{.*}} : !llvm<"<17 x i64>">
 //
 // 5. Rewrite as a masked write.
 //       CHECK: llvm.intr.masked.store %[[loaded]], %[[vecPtr_b]], %[[mask_b]]
-//  CHECK-SAME: {alignment = 128 : i32} :
+//  CHECK-SAME: {alignment = 4 : i32} :
 //  CHECK-SAME: !llvm<"<17 x float>">, !llvm<"<17 x i1>"> into !llvm<"<17 x float>*">
 
 func @transfer_read_2d_to_1d(%A : memref<?x?xf32>, %base0: index, %base1: index) -> vector<17xf32> {
   %f7 = constant 7.0: f32
   %f = vector.transfer_read %A[%base0, %base1], %f7
       {permutation_map = affine_map<(d0, d1) -> (d1)>} :
     memref<?x?xf32>, vector<17xf32>
   return %f: vector<17xf32>