[MLIR] Add std.assume_align op.

[rriddle]

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

2504

Are these patterns sorted alphabetically/can they be?

2533

nit: Move some of these to temporary variables. The nesting is a bit awkward.

2659

This breaks the ordering.

mlir/lib/Dialect/StandardOps/Ops.cpp

2772

Please use the declarative assembly form instead:
https://mlir.llvm.org/docs/OpDefinitions/#declarative-assembly-format

Should be something like:

let assemblyFormat = "$memref attr-dict `:` type($memref)";

2789

Why is the alignment not specified as an argument in ODS? That would autogenerate a lot of this for you automatically; builders, verification, etc.

[ftynse]

Thanks for pushing on this. Having alignment attribute as an ODS argument looks like it would remove a bunch of boilerplate for you.

mlir/include/mlir/Dialect/StandardOps/Ops.td

1654

Can you rather add alignment attribute as ODS argument, i.e. (ins AnyMemRef:$memref, I32Attr align) ? This will remove the need for a custom builder, and give you an auto-generated align() instead of hand-written getAlignment().

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

2504

They are not ordered, AFAIK. We should do it as a separate cleanup.

mlir/lib/Dialect/StandardOps/Ops.cpp

2772

Have you considered a non-dictionary syntax: assume_align 64, %memref ? This should be also possible with declarative assembly.

2789

I made the same comment above, +1!

[timshen]

Address all comments.

[timshen]

Done. Used ODS and assemblyFormat syntax. My brain didn't wire up to the simple and new stuff. Thanks!

[ftynse]

Thanks!

My brain didn't wire up to the simple and new stuff.

This is also why we have code review, the project is too big for one person to keep track of all the new shiny features. (Except River maybe).

[timshen]

Pinging @rriddle , can you take another look? Thanks!

[rriddle]

Looks great thanks!

mlir/lib/Dialect/StandardOps/Ops.cpp

2773

llvm::isPowerOf2_32

?

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

$res = builder.CreateAtomicCmpXchg($ptr, $cmp, $val,
                    getLLVMAtomicOrdering($success_ordering),
                    getLLVMAtomicOrdering($failure_ordering));
   }];
   let parser = [{ return parseAtomicCmpXchgOp(parser, result); }];
   let printer = [{ printAtomicCmpXchgOp(p, *this); }];
   let verifier = "return ::verify(*this);";
 }
 
+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});
+  }];
+}
+
 #endif // LLVMIR_OPS

mlir/include/mlir/Dialect/StandardOps/Ops.td

def ZeroExtendIOp : Std_Op<"zexti", [NoSideEffect, SameOperandsAndResultShape]> {
   let parser = [{
     return impl::parseCastOp(parser, result);
   }];
   let printer = [{
     return printStandardCastOp(this->getOperation(), p);
   }];
 }
 
+def AssumeAlignmentOp : Std_Op<"assume_alignment"> {
+  let summary =
+      "assertion that gives alignment information to the input memref";
+  let description = [{
+    The assume alignment operation takes a memref and a integer of alignment
+    value, and internally annotates the buffer with the given alignment. If
+    the buffer isn't aligned to the given alignment, the behavior is undefined.
+
+    This operation doesn't affect the semantics of a correct program. It's for
+    optimization only, and the optimization is best-effort.
+  }];
+  let arguments = (ins AnyMemRef:$memref, PositiveI32Attr:$alignment);
+  let results = (outs);
+
+  let assemblyFormat = "$memref `,` $alignment attr-dict `:` type($memref)";
+}
+
 #endif // STANDARD_OPS

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

for (int i = viewMemRefType.getRank() - 1; i >= 0; --i) {
       nextSize = size;
     }
 
     rewriter.replaceOp(op, {targetMemRef});
     return matchSuccess();
   }
 };
 
+struct AssumeAlignmentOpLowering
+    : public LLVMLegalizationPattern<AssumeAlignmentOp> {
+  using LLVMLegalizationPattern<AssumeAlignmentOp>::LLVMLegalizationPattern;
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    OperandAdaptor<AssumeAlignmentOp> transformed(operands);
+    Value memref = transformed.memref();
+    unsigned alignment = cast<AssumeAlignmentOp>(op).alignment().getZExtValue();
+
+    MemRefDescriptor memRefDescriptor(memref);
+    Value ptr = memRefDescriptor.alignedPtr(rewriter, memref.getLoc());
+
+    // Emit llvm.assume(memref.alignedPtr & (alignment - 1) == 0). Notice that
+    // the asserted memref.alignedPtr isn't used anywhere else, as the real
+    // users like load/store/views always re-extract memref.alignedPtr as they
+    // get lowered.
+    //
+    // This relies on LLVM's CSE optimization (potentially after SROA), since
+    // after CSE all memref.alignedPtr instances get de-duplicated into the same
+    // pointer SSA value.
+    Value zero =
+        createIndexAttrConstant(rewriter, op->getLoc(), getIndexType(), 0);
+    Value mask = createIndexAttrConstant(rewriter, op->getLoc(), getIndexType(),
+                                         alignment - 1);
+    Value ptrValue =
+        rewriter.create<LLVM::PtrToIntOp>(op->getLoc(), getIndexType(), ptr);
+    rewriter.create<LLVM::AssumeOp>(
+        op->getLoc(),
+        rewriter.create<LLVM::ICmpOp>(
+            op->getLoc(), LLVM::ICmpPredicate::eq,
+            rewriter.create<LLVM::AndOp>(op->getLoc(), ptrValue, mask), zero));
+
+    rewriter.eraseOp(op);
+    return matchSuccess();
+  }
+};
+
 } // namespace
 
 static void ensureDistinctSuccessors(Block &bb) {
   auto *terminator = bb.getTerminator();
 
   // Find repeated successors with arguments.
   llvm::SmallDenseMap<Block *, SmallVector<int, 4>> successorPositions;
   for (int i = 0, e = terminator->getNumSuccessors(); i < e; ++i) {
void mlir::populateStdToLLVMNonMemoryConversionPatterns(
   // clang-format on
 }
 
 void mlir::populateStdToLLVMMemoryConversionPatters(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     bool useAlloca) {
   // clang-format off
   patterns.insert<
+      AssumeAlignmentOpLowering,
       DimOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       StoreOpLowering,
       SubViewOpLowering,
       ViewOpLowering>(*converter.getDialect(), converter);
   patterns.insert<
       AllocOpLowering,
       DeallocOpLowering>(
         *converter.getDialect(), converter, useAlloca);
   // clang-format on
 }
 

mlir/lib/Dialect/StandardOps/Ops.cpp

for (auto resultStride : enumerate(resultStrides)) {
     if (resultStride.value() < baseStride ||
         resultStride.value() % baseStride != 0)
       return failure();
     staticStrides[resultStride.index()] = resultStride.value() / baseStride;
   }
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// AssumeAlignmentOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(AssumeAlignmentOp op) {
+  unsigned alignment = op.alignment().getZExtValue();
+  if (!llvm::isPowerOf2_32(alignment))
+    return op.emitOpError("alignment must be power of 2");
+  return success();
+}
+
 namespace {
 
 /// Pattern to rewrite a subview op with constant size arguments.
 class SubViewOpShapeFolder final : public OpRewritePattern<SubViewOp> {
 public:
   using OpRewritePattern<SubViewOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(SubViewOp subViewOp,
                                      PatternRewriter &rewriter) const override {
     MemRefType subViewType = subViewOp.getType();
     // Follow all or nothing approach for shapes for now. If all the operands
     // for sizes are constants then fold it into the type of the result memref.
     if (subViewType.hasStaticShape() ||
         llvm::any_of(subViewOp.sizes(), [](Value operand) {
           return !matchPattern(operand, m_ConstantIndex());
         })) {
       return matchFailure();
     }
     SmallVector<int64_t, 4> staticShape(subViewOp.getNumSizes());
     for (auto size : llvm::enumerate(subViewOp.sizes())) {

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

func @check_tanh_func_added_only_once_to_symbol_table(%f: f32, %lf: f64) -> () {
     %lf1 = std.tanh %lf0 : f64
     return
   }
 // CHECK: module {
 // CHECK: llvm.func @tanh(!llvm.double) -> !llvm.double
 // CHECK: llvm.func @tanhf(!llvm.float) -> !llvm.float
 // CHECK-LABEL: func @check_tanh_func_added_only_once_to_symbol_table
 }
+
+// -----
+
+// CHECK-LABEL: func @assume_alignment
+func @assume_alignment(%0 : memref<4x4xf16>) {
+  // CHECK: %[[PTR:.*]] = llvm.extractvalue %[[MEMREF:.*]][1] : !llvm<"{ half*, half*, i64, [2 x i64], [2 x i64] }">
+  // CHECK-NEXT: %[[ZERO:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[MASK:.*]] = llvm.mlir.constant(15 : index) : !llvm.i64
+  // CHECK-NEXT: %[[INT:.*]] = llvm.ptrtoint %[[PTR]] : !llvm<"half*"> to !llvm.i64
+  // CHECK-NEXT: %[[MASKED_PTR:.*]] = llvm.and %[[INT]], %[[MASK:.*]] : !llvm.i64
+  // CHECK-NEXT: %[[CONDITION:.*]] = llvm.icmp "eq" %[[MASKED_PTR]], %[[ZERO]] : !llvm.i64
+  // CHECK-NEXT: "llvm.intr.assume"(%[[CONDITION]]) : (!llvm.i1) -> ()
+  assume_alignment %0, 16 : memref<4x4xf16>
+  return
+}

mlir/test/IR/core-ops.mlir

 // CHECK-LABEL: func @tensor_load_store
 func @tensor_load_store(%0 : memref<4x4xi32>) {
   // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF:.*]] : memref<4x4xi32>
   %1 = tensor_load %0 : memref<4x4xi32>
   // CHECK: tensor_store %[[TENSOR]], %[[MEMREF]] : memref<4x4xi32>
   tensor_store %1, %0 : memref<4x4xi32>
   return
 }
+
+// CHECK-LABEL: func @assume_alignment
+// CHECK-SAME: %[[MEMREF:.*]]: memref<4x4xf16>
+func @assume_alignment(%0: memref<4x4xf16>) {
+  // CHECK: assume_alignment %[[MEMREF]], 16 : memref<4x4xf16>
+  assume_alignment %0, 16 : memref<4x4xf16>
+  return
+}

mlir/test/IR/invalid-ops.mlir

 // unranked to unranked
 func @invalid_memref_cast() {
   %0 = alloc() : memref<2x5xf32, 0>
   %1 = memref_cast %0 : memref<2x5xf32, 0> to memref<*xf32, 0>
   // expected-error@+1 {{operand type 'memref<*xf32>' and result type 'memref<*xf32>' are cast incompatible}}
   %2 = memref_cast %1 : memref<*xf32, 0> to memref<*xf32, 0>
   return
 }
+
+// -----
+
+// alignment is not power of 2.
+func @assume_alignment(%0: memref<4x4xf16>) {
+  // expected-error@+1 {{alignment must be power of 2}}
+  std.assume_alignment %0, 12 : memref<4x4xf16>
+  return
+}
+
+// -----
+
+// 0 alignment value.
+func @assume_alignment(%0: memref<4x4xf16>) {
+  // expected-error@+1 {{'std.assume_alignment' op attribute 'alignment' failed to satisfy constraint: positive 32-bit integer attribute}}
+  std.assume_alignment %0, 0 : memref<4x4xf16>
+  return
+}