[mlir][Memref] Introduce a memref::ExtractAlignedPointerAsIndexOp

As experience with memref::ExtractStridedMetadataOp grows we are
still missing a simple way to extract the pointer held by a memref
and lower to different backednds (LLVM, SPIRV, library calls).

This revision introduces a memref.extract_aligned_pointer_as_index that
returns an index containing the aligned pointer of the strided memref.

This operation is intended to be used solely as step during lowering,
it has no side effects. A reverse operation that creates a memref from
an index interpreted as a pointer is explicitly discouraged.

Differential Revision: https://reviews.llvm.org/D134651

diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index d223fb5..7445451 100644 (file)
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -787,11 +787,49 @@ def MemRef_DmaWaitOp : MemRef_Op<"dma_wait"> {
 }
 
 //===----------------------------------------------------------------------===//
-// ExtractMetadataOp
+// ExtractAlignedPointerAsIndexOp
+//===----------------------------------------------------------------------===//
+
+def MemRef_ExtractAlignedPointerAsIndexOp : MemRef_Op<"extract_aligned_pointer_as_index",
+    [NoSideEffect, SameVariadicResultSize]> {
+  let summary = "Extracts a memref's underlying aligned pointer as an index";
+  let description = [{
+    Extracts the underlying aligned pointer as an index.
+
+    This operation is useful for lowering to lower-level dialects while still
+    avoiding the need to define a pointer type in higher-level dialects such as
+    the memref dialect.
+
+    This operation is intended solely as step during lowering, it has no side
+    effects. A reverse operation that creates a memref from an index interpreted
+    as a pointer is explicitly discouraged.
+
+    Example:
+
+    ```
+      %0 = memref.extract_aligned_pointer_as_index %arg : memref<4x4xf32> -> index
+      %1 = arith.index_cast %0 : index to i64
+      %2 = llvm.inttoptr %1 : i64 to !llvm.ptr<f32>
+      call @foo(%2) : (!llvm.ptr<f32>) ->()
+    ```
+  }];
+
+  let arguments = (ins
+    AnyStridedMemRef:$source
+  );
+  let results = (outs Index:$aligned_pointer);
+
+  let assemblyFormat = [{
+    $source `:` type($source) `->` type(results) attr-dict
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// ExtractStridedMetadataOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_ExtractStridedMetadataOp : MemRef_Op<"extract_strided_metadata",
-    [SameVariadicResultSize]> {
+    [NoSideEffect, SameVariadicResultSize]> {
   let summary = "Extracts a buffer base with offset and strides";
   let description = [{
     Extracts a base buffer, offset and strides. This op allows additional layers

diff --git a/mlir/test/Dialect/MemRef/ops.mlir b/mlir/test/Dialect/MemRef/ops.mlir
index ca64a6a..b389c6f 100644 (file)
--- a/mlir/test/Dialect/MemRef/ops.mlir
+++ b/mlir/test/Dialect/MemRef/ops.mlir
@@ -374,3 +374,9 @@ func.func @memref_realloc_dd(%src : memref<?xf32>, %d: index)
   %0 = memref.realloc %src(%d) : memref<?xf32> to memref<?xf32>
   return %0 : memref<?xf32>
 }
+
+// CHECK-LABEL: func @memref_extract_aligned_pointer
+func.func @memref_extract_aligned_pointer(%src : memref<?xf32>) -> index {
+  %0 = memref.extract_aligned_pointer_as_index %src : memref<?xf32> -> index
+  return %0 : index
+}