[NFC][mlir][MemRef] Make use of InferTypeOpInterface

[qcolombet]

Summary

The InferTypeOpInterface generates builders for things it can infer the types.

Thanks to that interface we can:

• Eliminate a builder for DimOp, and
• Describe how to infer the result types of extract_strided_metadata from its source, and get a simpler builder as a result

NFC

[jpienaar]

This seems like implementing the type inference trait would also result in this.

[qcolombet]

This seems like implementing the type inference trait would also result in this.

I knew it was useful to do a review :).

Thanks @jpienaar !

Let me look at that closer.

[nicolasvasilache]

In D135734#3851289, @jpienaar wrote:

This seems like implementing the type inference trait would also result in this.

Does implementing the type inference trait result in automatically emitted builders via Tablegen?
If so, could we have some doc for this ?

Generally, it would be good to have an Interface section in the docs mirrored on the MLIR website like we have for dialects so that a high-level overview gives the key information before digging into particulars.

[nicolasvasilache]

I am ok with landing as is to get the stack in and implementing Jacques' suggestion as a followup

[qcolombet]

• Use InferTypeOpInterface

[qcolombet]

@jpienaar Could you take another look please?

[jpienaar]

LG thanks!

In D135734#3855171, @nicolasvasilache wrote:

In D135734#3851289, @jpienaar wrote:

This seems like implementing the type inference trait would also result in this.

Does implementing the type inference trait result in automatically emitted builders via Tablegen?
If so, could we have some doc for this ?

It does yes, and let me follow up on docs (I thought it was but I'll check in morning).

Generally, it would be good to have an Interface section in the docs mirrored on the MLIR website like we have for dialects so that a high-level overview gives the key information before digging into particulars.

Could you file an issue for this? This is good idea and we can autogenerate a lot of them ... We need to update a few things doc side (and definitely something which needs help and hands :-)).

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1257

You can use the OpAdaptor generated for this op if you want to use the generated accessor methods.

1262

It's been some time since I've worked with memrefs. To double check: they always have a memory space and rank? And it's fine to ignore layout here and just insert empty?

[qcolombet]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1257

Thanks for the pointer, let me double check that.

1262

Thanks for double checking.

For MemRefTypes, yes they always have a rank and memory space.

UnrankedMemRefType don't have a rank, but they inherit from BaseMemRefType, not MemRefType (i.e., MemRefType == ranked, UnrankedMemRefType == unranked). I believe this is how we tell them apart.

Regarding the layout, yes, this is what we want. Essentially this operation generate 0-D memref from an arbitrary memref.

[qcolombet]

• Use OpAdaptor to avoid using hardcoded indices

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 #define MLIR_DIALECT_MEMREF_IR_MEMREF_H_
 
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/ShapedOpInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 
 namespace arith {
 enum class AtomicRMWKind : uint64_t;

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

 //===- MemRefOps.td - MemRef op definitions ----------------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MEMREF_OPS
 #define MEMREF_OPS
 
 include "mlir/Dialect/Arith/IR/ArithBase.td"
 include "mlir/Dialect/MemRef/IR/MemRefBase.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/ShapedOpInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 
 /// A TypeAttr for memref types.
 def MemRefTypeAttr
def MemRef_DimOp : MemRef_Op<"dim", [
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
     attr-dict $source `,` $index `:` type($source)
   }];
 
   let builders = [
     OpBuilder<(ins "Value":$source, "int64_t":$index)>,
-    OpBuilder<(ins "Value":$source, "Value":$index)>
   ];
 
   let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     Optional<int64_t> getConstantIndex();
 
     /// Interface method of ShapedDimOpInterface: Return the source memref.
     Value getShapedValue() { return getSource(); }

 
 //===----------------------------------------------------------------------===//
 // ExtractStridedMetadataOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_ExtractStridedMetadataOp : MemRef_Op<"extract_strided_metadata", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     Pure, 
-    SameVariadicResultSize]> {
+    SameVariadicResultSize,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   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.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 
 void DimOp::build(OpBuilder &builder, OperationState &result, Value source,
                   int64_t index) {
   auto loc = result.location;
   Value indexValue = builder.create<arith::ConstantIndexOp>(loc, index);
   build(builder, result, source, indexValue);
 }
 
-void DimOp::build(OpBuilder &builder, OperationState &result, Value source,
-                  Value index) {
-  auto indexTy = builder.getIndexType();
-  build(builder, result, indexTy, source, index);
-}
-
 Optional<int64_t> DimOp::getConstantIndex() {
   if (auto constantOp = getIndex().getDefiningOp<arith::ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 
 Speculation::Speculatability DimOp::getSpeculatability() {
   auto constantIndex = getConstantIndex();
void ExtractAlignedPointerAsIndexOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "intptr");
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractStridedMetadataOp
 //===----------------------------------------------------------------------===//
 
+/// The number and type of the results are inferred from the
+/// shape of the source.
+LogicalResult ExtractStridedMetadataOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  ExtractStridedMetadataOpAdaptor extractAdaptor(operands, attributes, regions);
+  auto sourceType = extractAdaptor.getSource().getType().dyn_cast<MemRefType>();
+  if (!sourceType)
+    return failure();
+
+  unsigned sourceRank = sourceType.getRank();
+  IndexType indexType = IndexType::get(context);
+  auto memrefType =
+      MemRefType::get({}, sourceType.getElementType(),
+                      MemRefLayoutAttrInterface{}, sourceType.getMemorySpace());
+  // Base.
+  inferredReturnTypes.push_back(memrefType);
+  // Offset.
+  inferredReturnTypes.push_back(indexType);
+  // Sizes and strides.
+  for (unsigned i = 0; i < sourceRank * 2; ++i)
+    inferredReturnTypes.push_back(indexType);
+  return success();
+}
+
 void ExtractStridedMetadataOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getBaseBuffer(), "base_buffer");
   setNameFn(getOffset(), "offset");
   // For multi-result to work properly with pretty names and packed syntax `x:3`
   // we can only give a pretty name to the first value in the pack.
   if (!getSizes().empty()) {
     setNameFn(getSizes().front(), "sizes");

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
                                 PatternRewriter &rewriter) const override {
     auto subview = op.getSource().getDefiningOp<memref::SubViewOp>();
     if (!subview)
       return failure();
 
     // Build a plain extract_strided_metadata(memref) from
     // extract_strided_metadata(subview(memref)).
     Location origLoc = op.getLoc();
-    IndexType indexType = rewriter.getIndexType();
     Value source = subview.getSource();
     auto sourceType = source.getType().cast<MemRefType>();
     unsigned sourceRank = sourceType.getRank();
-    SmallVector<Type> sizeStrideTypes(sourceRank, indexType);
 
     auto newExtractStridedMetadata =
-        rewriter.create<memref::ExtractStridedMetadataOp>(
-            origLoc, op.getBaseBuffer().getType(), indexType, sizeStrideTypes,
-            sizeStrideTypes, source);
+        rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
 
     SmallVector<int64_t> sourceStrides;
     int64_t sourceOffset;
 
     bool hasKnownStridesAndOffset =
         succeeded(getStridesAndOffset(sourceType, sourceStrides, sourceOffset));
     (void)hasKnownStridesAndOffset;
     assert(hasKnownStridesAndOffset &&
LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
                                 PatternRewriter &rewriter) const override {
     auto reshape = op.getSource().getDefiningOp<ReassociativeReshapeLikeOp>();
     if (!reshape)
       return failure();
 
     // Build a plain extract_strided_metadata(memref) from
     // extract_strided_metadata(reassociative_reshape_like(memref)).
     Location origLoc = op.getLoc();
-    IndexType indexType = rewriter.getIndexType();
     Value source = reshape.getSrc();
     auto sourceType = source.getType().cast<MemRefType>();
     unsigned sourceRank = sourceType.getRank();
-    SmallVector<Type> sizeStrideTypes(sourceRank, indexType);
 
     auto newExtractStridedMetadata =
-        rewriter.create<memref::ExtractStridedMetadataOp>(
-            origLoc, op.getBaseBuffer().getType(), indexType, sizeStrideTypes,
-            sizeStrideTypes, source);
+        rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
 
     // Collect statically known information.
     SmallVector<int64_t> strides;
     int64_t offset;
     bool hasKnownStridesAndOffset =
         succeeded(getStridesAndOffset(sourceType, strides, offset));
     (void)hasKnownStridesAndOffset;
     assert(hasKnownStridesAndOffset &&