[mlir][MemRef] Introduce a memref.extract_metadata op.

[nicolasvasilache]

Summary

This is the counterpart of memref.reinterpret_cast and is useful to lift
strided memref manipulation out of the LLVM dialect.

Discussion: https://discourse.llvm.org/t/extracting-dynamic-offsets-strides-from-memref/64170

[stellaraccident]

This looks good to me but I'd like to see the reviews of the other contributors to the conversation before landing.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

705

Nit: it seems that it can be used anywhere (early or late) but it is resolved late?

[nicolasvasilache]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

705

Could you please elaborate what "resolved" means ?
I'm thinking it will fold away in various ways or continue to get materialized when going to LLVM (if uses remain).

[stellaraccident]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

705

My word "resolved" is likely the same as your word "folded" (but per the discussion on discourse, I think that such folding needs to be explicit as it has phase ordering constraints).

I was focusing more on "used": there really is not a constraint on where it gets used in various settings.

[Hardcode84]

Op itself is LGTM but maybe we should also add some simple test in ops.mlir?

Thanks

[bondhugula]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

705

Queries -> Extracts

705

I'd not mix commentary on usages in the pipeline early or late in the first para. Please move this to later if it all it's retained.

724–725

This can be brought up instead.

730

Missing memref..

757

Test cases are missing.

[bondhugula]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

703

Recommend avoiding "late/early" in the summary - that's all application-specific and relative. Instead,

Extracts a buffer base with offset and strides

?

705

Behavior of the op when such stride info can't be extracted has been left undocumented here.

730

extract_metadata would be too generic a name and not reflect what this op is returning. Instead, extract_stride_metadata or extract_stride_info looks appropriate.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

703

SGTM, thanks!

705

Opted for removing this as @bondhugula suggested.

I think the intended use case section is detailed enough.

730

went for memref.extract_strided_metadata

[bondhugula]

LGTM - but some comments to rephrase/rewrite certain things.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

708–709

716–728

I think these paras need some rephrasing - the first bullet refers to LLVM lowering, but the header talks about ops needing to manipulate stride info outside the LLVM dialect. I think this is a bullet on the latter is missing.

Also, I think "... to lift the materialization of the internal logic of ops that manipulate metadata" should be rewritten better.

Beyond these, I see this extraction needed for completeness sake just like we have dim op: you can have a memref passed to a function and extracting strides, sizes etc. may be needed to subsequently say construct a new memref -- let's say with slightly different strides, sizes but based on the previous memref.

717

metadata -> memref metadata?

725

"next problem" is unclear here.

733

Syntax broken here: sizes:2, ...

[nicolasvasilache]

Address

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

716–728

Rephrased and reshuffled a bit more, feel free to make another unreviewed pass over this and just land something that is more to your liking if this does not yet fit the bill.

Thanks!

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

unsigned getTagMemRefRank() {
       return getTagMemRef().getType().cast<MemRefType>().getRank();
     }
   }];
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// ExtractMetadataOp
+//===----------------------------------------------------------------------===//
+
+def MemRef_ExtractStridedMetadataOp : MemRef_Op<"extract_strided_metadata",
+    [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
+    of transformations and foldings to be added as lowering progresses from
+    higher-level dialect to lower-level dialects such as the LLVM dialect.
+
+    The op requires a strided memref source operand. If the source operand is not
+    a strided memref, then verification fails.
+
+    This operation is also useful for completeness to the existing memref.dim op.
+    While accessing strides, offsets and the base pointer independently is not
+    available, this is useful for composing with its natural complement op: 
+    `memref.reinterpret_cast`.
+
+    Intended Use Cases:
+
+    The main use case is to expose the logic for manipulate memref metadata at a
+    higher level than the LLVM dialect. 
+    This makes lowering more progressive and brings the following benefits:
+      - not all users of MLIR want to lower to LLVM and the information to e.g.
+        lower to library calls---like libxsmm---or to SPIR-V was not available.
+      - foldings and canonicalizations can happen at a higher level in MLIR: 
+        before this op existed, lowering to LLVM would create large amounts of 
+        LLVMIR. Even when LLVM does a good job at folding the low-level IR from
+        a performance perspective, it is unnecessarily opaque and inefficient to
+        send unkempt IR to LLVM.
+
+    Example:
+
+    ```mlir
+      %base, %offset, %sizes:2, %strides:2 = 
+        memref.extract_strided_metadata %memref : 
+          memref<10x?xf32>, index, index, index, index, index
+
+      // After folding, the type of %m2 can be memref<10x?xf32> and further 
+      // folded to %memref.
+      %m2 = memref.reinterpret_cast %base to
+          offset: [%offset],
+          sizes: [%sizes#0, %sizes#1],
+          strides: [%strides#0, %strides#1]
+        : memref<f32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+    ```
+  }];
+
+  let arguments = (ins
+    AnyStridedMemRef:$source
+  );
+  let results = (outs
+    AnyStridedMemRefOfRank<0>:$base_buffer,
+    Index:$offset,
+    Variadic<Index>:$sizes,
+    Variadic<Index>:$strides
+  );
+
+  let assemblyFormat = [{
+    $source `:` type($source) `->` type(results) attr-dict
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // GenericAtomicRMWOp
 //===----------------------------------------------------------------------===//
 
 def GenericAtomicRMWOp : MemRef_Op<"generic_atomic_rmw", [
       SingleBlockImplicitTerminator<"AtomicYieldOp">,
       TypesMatchWith<"result type matches element type of memref",
                      "memref", "result",
                      "$_self.cast<MemRefType>().getElementType()">

mlir/test/Dialect/MemRef/ops.mlir

// CHECK-NEXT: memref.generic_atomic_rmw [[BUF]]{{\[}}[[I]], [[J]]] : memref
     ^bb0(%old_value : f32):
       %c1 = arith.constant 1.0 : f32
       %out = arith.addf %c1, %old_value : f32
       memref.atomic_yield %out : f32
   // CHECK: index_attr = 8 : index
   } { index_attr = 8 : index }
   return
 }
+
+// -----
+
+func.func @extract_strided_metadata(%memref : memref<10x?xf32>) 
+    -> memref<?x?xf32, offset: ?, strides: [?, ?]> {
+
+  %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %memref
+    : memref<10x?xf32> -> memref<f32>, index, index, index, index, index
+
+  %m2 = memref.reinterpret_cast %base to
+      offset: [%offset],
+      sizes: [%sizes#0, %sizes#1],
+      strides: [%strides#0, %strides#1]
+    : memref<f32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+
+  return %m2: memref<?x?xf32, offset: ?, strides: [?, ?]>
+}