[MemRefToLLVM] Fix the lowering of memref.assume_alignment

[qcolombet]

Summary

memref.assume_alignment annotates the alignment of the source buffer
not the base pointer.
Put differently, prior to this patch memref.assume_alignment would lower
to llvm.assume %buffer.base.isAligned(X) whereas what we want is
llvm.assume (%buffer.base + %buffer.offset).isAligned(X).
In other words, we were missing to include the offset in the expression
checked by the llvm.assume.

[qcolombet]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

372

I was talking with @mravishankar offline and he pointed out that if we are aligned at (alignPtr + offset) then alignPtr must be aligned at least with the same alignment.
I.e., we could issue two assumes here:

• one for the aligned base pointer.
• one for the aligned base pointer + offset.

It sounded like useful information to lower.

What do you guys think?

[mravishankar]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

372

To clarify this a bit. If you need an basePtr + offset to be aligned to a value v, you need both basePtr and offset to be aligned to v. Without that you cannot guarantee that basePtr + offset is always aligned to v for all values of basePtr.
This is the reverse logic though. Here it is saying basePtr + offset is aligned, but I am not sure if can always assume that basePtr is aligned. Maybe we skip it for now... In most common case though, basePtr is also aligned. We can let the user (in this case IREE) generate assume for both.

[mravishankar]

Looks good to me. It is strange that `memref.assume_alignment %0, 16 : memref<4x4xf16, strided<[?, ?], offset: ?>>` refers to the pointer after offset, and not the base. But I understand the reasoning here.

[qcolombet]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

372

Thinking about it a little more, I think it is better to let the producers of memref.assume_alignment to separately generate one for the base pointer if they want.
I can think of counter examples where this wouldn't be correct to issue that assume.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

struct AssumeAlignmentOpLowering
 
   LogicalResult
   matchAndRewrite(memref::AssumeAlignmentOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Value memref = adaptor.getMemref();
     unsigned alignment = op.getAlignment();
     auto loc = op.getLoc();
 
+    auto srcMemRefType = op.getMemref().getType().cast<MemRefType>();
+    // When we convert to LLVM, the input memref must have been normalized
+    // beforehand. Hence, this call is guaranteed to work.
+    auto [strides, offset] = getStridesAndOffset(srcMemRefType);
+
     MemRefDescriptor memRefDescriptor(memref);
     Value ptr = memRefDescriptor.alignedPtr(rewriter, memref.getLoc());
+    // Skip if offset is zero.
+    if (offset != 0) {
+      Value offsetVal = ShapedType::isDynamic(offset)
+                            ? memRefDescriptor.offset(rewriter, loc)
+                            : createIndexConstant(rewriter, loc, offset);
+      Type elementType =
+          typeConverter->convertType(srcMemRefType.getElementType());
+      ptr = rewriter.create<LLVM::GEPOp>(loc, ptr.getType(), elementType, ptr,
+                                         offsetVal);
+    }
 
-    // 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.
+    // Emit llvm.assume(memref & (alignment - 1) == 0).
     //
     // This relies on LLVM's CSE optimization (potentially after SROA), since
-    // after CSE all memref.alignedPtr instances get de-duplicated into the same
+    // after CSE all memref instances should get de-duplicated into the same
     // pointer SSA value.
     auto intPtrType =
         getIntPtrType(memRefDescriptor.getElementPtrType().getAddressSpace());
     Value zero = createIndexAttrConstant(rewriter, loc, intPtrType, 0);
     Value mask =
         createIndexAttrConstant(rewriter, loc, intPtrType, alignment - 1);
     Value ptrValue = rewriter.create<LLVM::PtrToIntOp>(loc, intPtrType, ptr);
     rewriter.create<LLVM::AssumeOp>(
         loc, rewriter.create<LLVM::ICmpOp>(
                  loc, LLVM::ICmpPredicate::eq,
                  rewriter.create<LLVM::AndOp>(loc, ptrValue, mask), zero));
 
     rewriter.eraseOp(op);
     return success();
   }
 };

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

func.func @subview(%0 : memref<64x4xf32, strided<[4, 1], offset: 0>>, %arg0 : index, %arg1 : index, %arg2 : index) {
   %1 = memref.subview %0[%arg0, %arg1][%arg0, %arg1][%arg0, %arg1] :
     memref<64x4xf32, strided<[4, 1], offset: 0>>
   to memref<?x?xf32, strided<[?, ?], offset: ?>>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @assume_alignment
+// CHECK-LABEL: func @assume_alignment(
 func.func @assume_alignment(%0 : memref<4x4xf16>) {
   // CHECK: %[[PTR:.*]] = llvm.extractvalue %[[MEMREF:.*]][1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
   // CHECK-NEXT: %[[ZERO:.*]] = llvm.mlir.constant(0 : index) : i64
   // CHECK-NEXT: %[[MASK:.*]] = llvm.mlir.constant(15 : index) : i64
   // CHECK-NEXT: %[[INT:.*]] = llvm.ptrtoint %[[PTR]] : !llvm.ptr to i64
   // CHECK-NEXT: %[[MASKED_PTR:.*]] = llvm.and %[[INT]], %[[MASK:.*]] : i64
   // CHECK-NEXT: %[[CONDITION:.*]] = llvm.icmp "eq" %[[MASKED_PTR]], %[[ZERO]] : i64
   // CHECK-NEXT: "llvm.intr.assume"(%[[CONDITION]]) : (i1) -> ()
   memref.assume_alignment %0, 16 : memref<4x4xf16>
   return
 }
 
 // -----
 
+// CHECK-LABEL: func @assume_alignment_w_offset
+func.func @assume_alignment_w_offset(%0 : memref<4x4xf16, strided<[?, ?], offset: ?>>) {
+  // CHECK-DAG: %[[PTR:.*]] = llvm.extractvalue %[[MEMREF:.*]][1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK-DAG: %[[OFFSET:.*]] = llvm.extractvalue %[[MEMREF]][2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK-DAG: %[[BUFF_ADDR:.*]] =  llvm.getelementptr %[[PTR]][%[[OFFSET]]] : (!llvm.ptr, i64) -> !llvm.ptr, f16
+  // CHECK-DAG: %[[ZERO:.*]] = llvm.mlir.constant(0 : index) : i64
+  // CHECK-DAG: %[[MASK:.*]] = llvm.mlir.constant(15 : index) : i64
+  // CHECK-NEXT: %[[INT:.*]] = llvm.ptrtoint %[[BUFF_ADDR]] : !llvm.ptr to i64
+  // CHECK-NEXT: %[[MASKED_PTR:.*]] = llvm.and %[[INT]], %[[MASK:.*]] : i64
+  // CHECK-NEXT: %[[CONDITION:.*]] = llvm.icmp "eq" %[[MASKED_PTR]], %[[ZERO]] : i64
+  // CHECK-NEXT: "llvm.intr.assume"(%[[CONDITION]]) : (i1) -> ()
+  memref.assume_alignment %0, 16 : memref<4x4xf16, strided<[?, ?], offset: ?>>
+  return
+}
+// -----
+
 // CHECK-LABEL: func @dim_of_unranked
 // CHECK32-LABEL: func @dim_of_unranked
 func.func @dim_of_unranked(%unranked: memref<*xi32>) -> index {
   %c0 = arith.constant 0 : index
   %dim = memref.dim %unranked, %c0 : memref<*xi32>
   return %dim : index
 }
 // CHECK: %[[UNRANKED_DESC:.*]] = builtin.unrealized_conversion_cast