Fix how we get alignment for memrefs

Author: Lily Orth-Smith

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

@@ -70,15 +70,26 @@ 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)
+  // If shape is statically known, assign MemRefTypes to the alignment of a
+  // VectorType with the same size and dtype. Otherwise, fall back to the
+  // alignment of the element type.
+  Type convertedType;
+  if (memrefType.hasStaticShape()) {
+    convertedType = typeConverter.convertType(VectorType::get(
+        memrefType.getNumElements(), memrefType.getElementType()));
+  } else {
+    convertedType = typeConverter.convertType(memrefType.getElementType());
+  }
+
+  if (!convertedType)
     return failure();
 
   // TODO: this should use the MLIR data layout when it becomes available and
   // stop depending on translation.
   llvm::LLVMContext llvmContext;
-  align = LLVM::TypeToLLVMIRTranslator(llvmContext)
-              .getPreferredAlignment(elementTy, typeConverter.getDataLayout());
+  align =
+      LLVM::TypeToLLVMIRTranslator(llvmContext)
+          .getPreferredAlignment(convertedType, typeConverter.getDataLayout());
   return success();
 }