[mlir][memref]Add conversion support for memref.extract_aligned_pointer_as_index to LLVM

[ftynse]

I don't understand why the pattern is restricted to allocs and block arguments. One can trivially have a non-alloc operation than pass its result through a branch to make it a block argument.

[nicolasvasilache]

Address comment.

[nicolasvasilache]

In D134834#3823179, @ftynse wrote:

I don't understand why the pattern is restricted to allocs and block arguments. One can trivially have a non-alloc operation than pass its result through a branch to make it a block argument.

No good reason, I was orthogonalizing unnecessarily.

[nicolasvasilache]

Fix test

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

auto dataPtr =
                              adaptor.getIndices(), rewriter);
     rewriter.replaceOpWithNewOp<LLVM::AtomicRMWOp>(
         atomicOp, resultType, *maybeKind, dataPtr, adaptor.getValue(),
         LLVM::AtomicOrdering::acq_rel);
     return success();
   }
 };
 
+/// Unpack the pointer returned by a memref.extract_aligned_pointer_as_index.
+class ConvertExtractAlignedPointerAsIndex
+    : public ConvertOpToLLVMPattern<memref::ExtractAlignedPointerAsIndexOp> {
+public:
+  using ConvertOpToLLVMPattern<
+      memref::ExtractAlignedPointerAsIndexOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::ExtractAlignedPointerAsIndexOp extractOp,
+                  OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefDescriptor desc(adaptor.getSource());
+    rewriter.replaceOpWithNewOp<LLVM::PtrToIntOp>(
+        extractOp, getTypeConverter()->getIndexType(),
+        desc.alignedPtr(rewriter, extractOp->getLoc()));
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::populateMemRefToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
       AllocaOpLowering,
       AllocaScopeOpLowering,
       AtomicRMWOpLowering,
       AssumeAlignmentOpLowering,
+      ConvertExtractAlignedPointerAsIndex,
       DimOpLowering,
       GenericAtomicRMWOpLowering,
       GlobalMemrefOpLowering,
       GetGlobalMemrefOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       MemRefCopyOpLowering,
       MemRefReinterpretCastOpLowering,

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

func.func @memref_copy_unranked() {
   // CHECK: [[ALLOCA2:%.*]] = llvm.alloca [[RANK2]] x !llvm.struct<(i64, ptr<i8>)> : (i64) -> !llvm.ptr<struct<(i64, ptr<i8>)>>
   // CHECK: llvm.store {{%.*}}, [[ALLOCA2]] : !llvm.ptr<struct<(i64, ptr<i8>)>>
   // CHECK: [[ALLOCA3:%.*]] = llvm.alloca [[RANK2]] x !llvm.struct<(i64, ptr<i8>)> : (i64) -> !llvm.ptr<struct<(i64, ptr<i8>)>>
   // CHECK: llvm.store [[INSERT2]], [[ALLOCA3]] : !llvm.ptr<struct<(i64, ptr<i8>)>>
   // CHECK: [[SIZE:%.*]] = llvm.mlir.constant(1 : index) : i64
   // CHECK: llvm.call @memrefCopy([[SIZE]], [[ALLOCA2]], [[ALLOCA3]]) : (i64, !llvm.ptr<struct<(i64, ptr<i8>)>>, !llvm.ptr<struct<(i64, ptr<i8>)>>) -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @extract_aligned_pointer_as_index
+func.func @extract_aligned_pointer_as_index(%m: memref<?xf32>) -> index {
+  %0 = memref.extract_aligned_pointer_as_index %m: memref<?xf32> -> index
+  // CHECK: %[[E:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)> 
+  // CHECK: %[[I64:.*]] = llvm.ptrtoint %[[E]] : !llvm.ptr<f32> to i64
+  // CHECK: %[[R:.*]] = builtin.unrealized_conversion_cast %[[I64]] : i64 to index
+
+  // CHECK: return %[[R:.*]] : index
+  return %0: index
+}