[mlir][vector] set alignment when lowering transfer_read and transfer_write.

[whchung]

Summary

When emitting masked load / store, set alignment from data layout.

[whchung]

@ftynse This patch is somewhat related to our discussion in D79019. It now becomes evident it's better to specify target triple and data layout in llvm dialect.

[ftynse]

It now becomes evident it's better to specify target triple and data layout in llvm dialect.

It should be in the *->llvm conversion. However, the conversion is not very well structured for extension atm, and we would have to patch all *->llvm to take the data layout and forward it to the relevant palce.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

755

We tend to use arguments passed by-reference to return extra results, or named structs.

757

Type converter mail fail and return nullptr, what's the contract with this function on it?

[whchung]

Address code review comments adding more checks.

[whchung]

@ftynse Revised the patch addressing your comments. Could you help review it again, and help get it landed should the patch is okay with you? Thanks.

[nicolasvasilache]

LGTM, thanks @whchung

[aartbik]

This breaks transfer read/write operations when using the mlir cpu runner in JIT mode, probably due to inconsistencies between the data layouts.

[rriddle]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

755

nit: Static functions should be marked static and in the top-level namespace, i.e., not within an anonymous namespace.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 };
 
 template <typename ConcreteOp>
 void replaceTransferOp(ConversionPatternRewriter &rewriter,
                        LLVMTypeConverter &typeConverter, Location loc,
                        Operation *op, ArrayRef<Value> operands, Value dataPtr,
                        Value mask);
 
+LogicalResult getLLVMTypeAndAlignment(LLVMTypeConverter &typeConverter,
+                                      Type type, LLVM::LLVMType &llvmType,
+                                      unsigned &align) {
+  auto convertedType = typeConverter.convertType(type);
+  if (!convertedType)
+    return failure();
+
+  llvmType = convertedType.template cast<LLVM::LLVMType>();
+  auto dataLayout = typeConverter.getDialect()->getLLVMModule().getDataLayout();
+  align = dataLayout.getPrefTypeAlignment(llvmType.getUnderlyingType());
+  return success();
+}
+
 template <>
 void replaceTransferOp<TransferReadOp>(ConversionPatternRewriter &rewriter,
                                        LLVMTypeConverter &typeConverter,
                                        Location loc, Operation *op,
                                        ArrayRef<Value> operands, Value dataPtr,
                                        Value mask) {
   auto xferOp = cast<TransferReadOp>(op);
   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);
 
-  auto vecTy = toLLVMTy(xferOp.getVectorType()).template cast<LLVM::LLVMType>();
+  LLVM::LLVMType vecTy;
+  unsigned align;
+  if (succeeded(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
+                                        vecTy, align)))
     rewriter.replaceOpWithNewOp<LLVM::MaskedLoadOp>(
         op, vecTy, dataPtr, mask, ValueRange{fill},
-      rewriter.getI32IntegerAttr(1));
+        rewriter.getI32IntegerAttr(align));
 }
 
 template <>
 void replaceTransferOp<TransferWriteOp>(ConversionPatternRewriter &rewriter,
                                         LLVMTypeConverter &typeConverter,
                                         Location loc, Operation *op,
                                         ArrayRef<Value> operands, Value dataPtr,
                                         Value mask) {
   auto adaptor = TransferWriteOpOperandAdaptor(operands);
+
+  auto xferOp = cast<TransferWriteOp>(op);
+  LLVM::LLVMType vecTy;
+  unsigned align;
+  if (succeeded(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
+                                        vecTy, align)))
     rewriter.replaceOpWithNewOp<LLVM::MaskedStoreOp>(
-      op, adaptor.vector(), dataPtr, mask, rewriter.getI32IntegerAttr(1));
+        op, adaptor.vector(), dataPtr, mask, rewriter.getI32IntegerAttr(align));
 }
 
 static TransferReadOpOperandAdaptor
 getTransferOpAdapter(TransferReadOp xferOp, ArrayRef<Value> operands) {
   return TransferReadOpOperandAdaptor(operands);
 }
 
 static TransferWriteOpOperandAdaptor

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 = 1 : i32} :
+//  CHECK-SAME: %[[PASS_THROUGH]] {alignment = 128 : 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 = 1 : i32} :
+//  CHECK-SAME: {alignment = 128 : 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>