[mlir] Propagate alignment annotations to LLVM loads and stores.

[timshen]

Backport.

[timshen]

Fix an arc mistake.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: fail. 62144 tests passed, 7 failed and 811 were skipped.

failed: MLIR.Conversion/StandardToLLVM/convert-memref-ops.mlir
failed: MLIR.Transforms/memref-normalize.mlir
failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: fail. 62144 tests passed, 7 failed and 811 were skipped.

failed: MLIR.Conversion/StandardToLLVM/convert-memref-ops.mlir
failed: MLIR.Transforms/memref-normalize.mlir
failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

static MemRefDescriptor fromStaticShape(OpBuilder &builder, Location loc,
                                           MemRefType type, Value memory);
 
   /// Builds IR extracting the allocated pointer from the descriptor.
   Value allocatedPtr(OpBuilder &builder, Location loc);
   /// Builds IR inserting the allocated pointer into the descriptor.
   void setAllocatedPtr(OpBuilder &builder, Location loc, Value ptr);
 
   /// Builds IR extracting the aligned pointer from the descriptor.
-  Value alignedPtr(OpBuilder &builder, Location loc);
+  Value alignedPtr(OpBuilder &builder, unsigned alignment, Location loc);
 
   /// Builds IR inserting the aligned pointer into the descriptor.
   void setAlignedPtr(OpBuilder &builder, Location loc, Value ptr);
 
   /// Builds IR extracting the offset from the descriptor.
   Value offset(OpBuilder &builder, Location loc);
 
   /// Builds IR inserting the offset into the descriptor.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_Prefetch : LLVM_ZeroResultOp<"intr.prefetch">,
   let llvmBuilder = [{
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::Function *fn = llvm::Intrinsic::getDeclaration(
         module, llvm::Intrinsic::prefetch, $addr->getType());
     builder.CreateCall(fn, {$addr, $rw, $hint, $cache});
   }];
 }
 
+def LLVM_AssumeOp : LLVM_Op<"intr.assume", []>,
+                    Arguments<(ins LLVM_Type:$cond)> {
+  let llvmBuilder = [{
+    llvm::Module *module = builder.GetInsertBlock()->getModule();
+    llvm::Function *fn =
+        llvm::Intrinsic::getDeclaration(module, llvm::Intrinsic::assume, {});
+    builder.CreateCall(fn, {$cond});
+  }];
+}
+
 def AtomicBinOpXchg : I64EnumAttrCase<"xchg", 0>;
 def AtomicBinOpAdd  : I64EnumAttrCase<"add", 1>;
 def AtomicBinOpSub  : I64EnumAttrCase<"sub", 2>;
 def AtomicBinOpAnd  : I64EnumAttrCase<"_and", 3>;
 def AtomicBinOpNand : I64EnumAttrCase<"nand", 4>;
 def AtomicBinOpOr   : I64EnumAttrCase<"_or", 5>;
 def AtomicBinOpXor  : I64EnumAttrCase<"_xor", 6>;
 def AtomicBinOpMax  : I64EnumAttrCase<"max", 7>;

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

public:
   BaseViewConversionHelper(Type type)
       : d(MemRefDescriptor::undef(rewriter(), loc(), type)) {}
 
   BaseViewConversionHelper(Value v) : d(v) {}
 
   /// Wrappers around MemRefDescriptor that use EDSC builder and location.
   Value allocatedPtr() { return d.allocatedPtr(rewriter(), loc()); }
   void setAllocatedPtr(Value v) { d.setAllocatedPtr(rewriter(), loc(), v); }
-  Value alignedPtr() { return d.alignedPtr(rewriter(), loc()); }
+  Value alignedPtr() { return d.alignedPtr(rewriter(), 0, loc()); }
   void setAlignedPtr(Value v) { d.setAlignedPtr(rewriter(), loc(), v); }
   Value offset() { return d.offset(rewriter(), loc()); }
   void setOffset(Value v) { d.setOffset(rewriter(), loc(), v); }
   Value size(unsigned i) { return d.size(rewriter(), loc(), i); }
   void setSize(unsigned i, Value v) { d.setSize(rewriter(), loc(), i, v); }
   void setConstantSize(unsigned i, int64_t v) {
     d.setConstantSize(rewriter(), loc(), i, v);
   }

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

 }
 
 /// Builds IR inserting the allocated pointer into the descriptor.
 void MemRefDescriptor::setAllocatedPtr(OpBuilder &builder, Location loc,
                                        Value ptr) {
   setPtr(builder, loc, kAllocatedPtrPosInMemRefDescriptor, ptr);
 }
 
+// Creates a constant Op producing a value of `resultType` from an index-typed
+// integer attribute.
+static Value createIndexAttrConstant(OpBuilder &builder, Location loc,
+                                     Type resultType, int64_t value) {
+  return builder.create<LLVM::ConstantOp>(
+      loc, resultType, builder.getIntegerAttr(builder.getIndexType(), value));
+}
+
 /// Builds IR extracting the aligned pointer from the descriptor.
-Value MemRefDescriptor::alignedPtr(OpBuilder &builder, Location loc) {
-  return extractPtr(builder, loc, kAlignedPtrPosInMemRefDescriptor);
+Value MemRefDescriptor::alignedPtr(OpBuilder &builder, unsigned alignment,
+                                   Location loc) {
+  Value ptr = extractPtr(builder, loc, kAlignedPtrPosInMemRefDescriptor);
+  if (alignment) {
+    assert(((alignment - 1) & alignment) == 0 &&
+           "Alignments must be power of 2");
+    builder.create<LLVM::AssumeOp>(
+        loc, builder.create<LLVM::ICmpOp>(
+                 loc, LLVM::ICmpPredicate::eq,
+                 builder.create<LLVM::AndOp>(
+                     loc, builder.create<LLVM::PtrToIntOp>(loc, indexType, ptr),
+                     createIndexAttrConstant(builder, loc, indexType,
+                                             alignment - 1)),
+                 createIndexAttrConstant(builder, loc, indexType, 0)));
+  }
+  return ptr;
 }
 
 /// Builds IR inserting the aligned pointer into the descriptor.
 void MemRefDescriptor::setAlignedPtr(OpBuilder &builder, Location loc,
                                      Value ptr) {
   setPtr(builder, loc, kAlignedPtrPosInMemRefDescriptor, ptr);
 }
 
-// Creates a constant Op producing a value of `resultType` from an index-typed
-// integer attribute.
-static Value createIndexAttrConstant(OpBuilder &builder, Location loc,
-                                     Type resultType, int64_t value) {
-  return builder.create<LLVM::ConstantOp>(
-      loc, resultType, builder.getIntegerAttr(builder.getIndexType(), value));
-}
-
 /// Builds IR extracting the offset from the descriptor.
 Value MemRefDescriptor::offset(OpBuilder &builder, Location loc) {
   return builder.create<LLVM::ExtractValueOp>(
       loc, indexType, value,
       builder.getI64ArrayAttr(kOffsetPosInMemRefDescriptor));
 }
 
 /// Builds IR inserting the offset into the descriptor.
for (int i = 1, nSizes = allocSizes.size(); i < nSizes; ++i) {
           loc, this->getIndexType(), ArrayRef<Value>{linearized, indices[i]});
     }
     return linearized;
   }
 
   // This is a strided getElementPtr variant that linearizes subscripts as:
   //   `base_offset + index_0 * stride_0 + ... + index_n * stride_n`.
   Value getStridedElementPtr(Location loc, Type elementTypePtr,
-                             Value descriptor, ArrayRef<Value> indices,
-                             ArrayRef<int64_t> strides, int64_t offset,
+                             unsigned alignment, Value descriptor,
+                             ArrayRef<Value> indices, ArrayRef<int64_t> strides,
+                             int64_t offset,
                              ConversionPatternRewriter &rewriter) const {
     MemRefDescriptor memRefDescriptor(descriptor);
 
-    Value base = memRefDescriptor.alignedPtr(rewriter, loc);
+    Value base = memRefDescriptor.alignedPtr(rewriter, alignment, loc);
     Value offsetValue = offset == MemRefType::getDynamicStrideOrOffset()
                             ? memRefDescriptor.offset(rewriter, loc)
                             : this->createIndexConstant(rewriter, loc, offset);
 
     for (int i = 0, e = indices.size(); i < e; ++i) {
       Value stride = strides[i] == MemRefType::getDynamicStrideOrOffset()
                          ? memRefDescriptor.stride(rewriter, loc, i)
                          : this->createIndexConstant(rewriter, loc, strides[i]);
Value getDataPtr(Location loc, MemRefType type, Value memRefDesc,
                    ArrayRef<Value> indices, ConversionPatternRewriter &rewriter,
                    llvm::Module &module) const {
     LLVM::LLVMType ptrType = MemRefDescriptor(memRefDesc).getElementType();
     int64_t offset;
     SmallVector<int64_t, 4> strides;
     auto successStrides = getStridesAndOffset(type, strides, offset);
     assert(succeeded(successStrides) && "unexpected non-strided memref");
     (void)successStrides;
-    return getStridedElementPtr(loc, ptrType, memRefDesc, indices, strides,
-                                offset, rewriter);
+    return getStridedElementPtr(loc, ptrType, type.getAlignment(), memRefDesc,
+                                indices, strides, offset, rewriter);
   }
 };
 
 // Load operation is lowered to obtaining a pointer to the indexed element
 // and loading it.
 struct LoadOpLowering : public LoadStoreOpLowering<LoadOp> {
   using Base::Base;
 
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);
 
     // Copy the buffer pointer from the old descriptor to the new one.
     Value extracted = sourceMemRef.allocatedPtr(rewriter, loc);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
-    extracted = sourceMemRef.alignedPtr(rewriter, loc);
+    extracted = sourceMemRef.alignedPtr(rewriter, 0, loc);
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     // Extract strides needed to compute offset.
     SmallVector<Value, 4> strideValues;
     strideValues.reserve(viewMemRefType.getRank());
     for (int i = 0, e = viewMemRefType.getRank(); i < e; ++i)
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
 
     // Field 1: Copy the allocated pointer, used for malloc/free.
     Value extracted = sourceMemRef.allocatedPtr(rewriter, loc);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
     // Field 2: Copy the actual aligned pointer to payload.
-    extracted = sourceMemRef.alignedPtr(rewriter, loc);
+    extracted = sourceMemRef.alignedPtr(rewriter, 0, loc);
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     // Field 3: Copy the offset in aligned pointer.
     unsigned numDynamicSizes = llvm::size(viewOp.getDynamicSizes());
     (void)numDynamicSizes;
     bool hasDynamicOffset = offset == MemRefType::getDynamicStrideOrOffset();

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);
     Type llvmTargetElementTy = desc.getElementType();
     // Set allocated ptr.
     Value allocated = sourceMemRef.allocatedPtr(rewriter, loc);
     allocated =
         rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated);
     desc.setAllocatedPtr(rewriter, loc, allocated);
     // Set aligned ptr.
-    Value ptr = sourceMemRef.alignedPtr(rewriter, loc);
+    Value ptr = sourceMemRef.alignedPtr(rewriter, 0, loc);
     ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr);
     desc.setAlignedPtr(rewriter, loc, ptr);
     // Fill offset 0.
     auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0);
     auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr);
     desc.setOffset(rewriter, loc, zero);
 
     // Fill size and stride descriptors in memref.

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

// CHECK-NEXT:  llvm.mlir.constant(15 : index) : !llvm.i64
   %2 = dim %static, 2 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(13 : index) : !llvm.i64
   %3 = dim %static, 3 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(27 : index) : !llvm.i64
   %4 = dim %static, 4 : memref<42x32x15x13x27xf32>
   return
 }
 
+// CHECK-LABEL: func @aligned_memref_load
+//       CHECK:   %[[A:.*]]: !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">, %[[I:.*]]: !llvm.i64
+func @aligned_memref_load(%arr : memref<4xf16, 0, 8>, %index : index) -> f16 {
+  // CHECK-NEXT: %[[ld:.*]] = llvm.load %{{.*}} : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">
+  // CHECK-NEXT: %[[ptr:.*]] = llvm.extractvalue %[[ld]][1] : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }">
+
+  // CHECK-NEXT: %[[zero:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[mask:.*]] = llvm.mlir.constant(7 : index) : !llvm.i64
+  // CHECK-NEXT: %[[intptr:.*]] = llvm.ptrtoint %1 : !llvm<"half*"> to !llvm.i64
+  // CHECK-NEXT: %[[masked:.*]] = llvm.and %[[intptr]], %[[mask]] : !llvm.i64
+  // CHECK-NEXT: %[[isaligned:.*]] = llvm.icmp "eq" %[[masked]], %[[zero]] : !llvm.i64
+  // CHECK-NEXT: "llvm.intr.assume"(%[[isaligned]]) : (!llvm.i1) -> ()
+
+  // CHECK-NEXT: %[[off:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[st0:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+  // CHECK-NEXT: %[[offI:.*]] = llvm.mul %[[I]], %[[st0]] : !llvm.i64
+  // CHECK-NEXT: %[[off0:.*]] = llvm.add %[[off]], %[[offI]] : !llvm.i64
+  // CHECK-NEXT: %[[addr:.*]] = llvm.getelementptr %[[ptr]][%[[off0]]] : (!llvm<"half*">, !llvm.i64) -> !llvm<"half*">
+  // CHECK-NEXT: llvm.load %[[addr]] : !llvm<"half*">
+  %v = load %arr[%index] : memref<4xf16, 0, 8>
+  return %v : f16
+}
+
+// CHECK-LABEL: func @aligned_memref_store
+//       CHECK:   %[[A:.*]]: !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">, %[[V:.*]]: !llvm.half, %[[I:.*]]: !llvm.i64
+func @aligned_memref_store(%arr : memref<4xf16, 0, 8>, %v : f16, %index : index) {
+  // CHECK-NEXT: %[[ld:.*]] = llvm.load %{{.*}} : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">
+  // CHECK-NEXT: %[[ptr:.*]] = llvm.extractvalue %[[ld]][1] : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }">
+
+  // CHECK-NEXT: %[[zero:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[mask:.*]] = llvm.mlir.constant(7 : index) : !llvm.i64
+  // CHECK-NEXT: %[[intptr:.*]] = llvm.ptrtoint %1 : !llvm<"half*"> to !llvm.i64
+  // CHECK-NEXT: %[[masked:.*]] = llvm.and %[[intptr]], %[[mask]] : !llvm.i64
+  // CHECK-NEXT: %[[isaligned:.*]] = llvm.icmp "eq" %[[masked]], %[[zero]] : !llvm.i64
+  // CHECK-NEXT: "llvm.intr.assume"(%[[isaligned]]) : (!llvm.i1) -> ()
+
+  // CHECK-NEXT: %[[off:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[st0:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+  // CHECK-NEXT: %[[offI:.*]] = llvm.mul %[[I]], %[[st0]] : !llvm.i64
+  // CHECK-NEXT: %[[off0:.*]] = llvm.add %[[off]], %[[offI]] : !llvm.i64
+  // CHECK-NEXT: %[[addr:.*]] = llvm.getelementptr %[[ptr]][%[[off0]]] : (!llvm<"half*">, !llvm.i64) -> !llvm<"half*">
+  // CHECK-NEXT: llvm.store %[[V]], %[[addr]] : !llvm<"half*">
+  store %v, %arr[%index] : memref<4xf16, 0, 8>
+  return
+}