[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

1249

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

1254

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

1249

Thanks for the pointer, let me double check that.

1254

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/MemRefOps.td

def MemRef_ExtractStridedMetadataOp : MemRef_Op<"extract_strided_metadata", [
   );
   let results = (outs
     AnyStridedMemRefOfRank<0>:$base_buffer,
     Index:$offset,
     Variadic<Index>:$sizes,
     Variadic<Index>:$strides
   );
 
+  // Build `extract_strided_metadata(source)`.
+  // The number and type of the results are inferred from the
+  // shape of the source.
+  let builders = [OpBuilder<(ins "Value":$source)>];
+
   let assemblyFormat = [{
     $source `:` type($source) `->` type(results) attr-dict
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // GenericAtomicRMWOp
 //===----------------------------------------------------------------------===//

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

void ExtractAlignedPointerAsIndexOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "intptr");
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractStridedMetadataOp
 //===----------------------------------------------------------------------===//
 
+void ExtractStridedMetadataOp::build(mlir::OpBuilder &builder,
+                                     mlir::OperationState &state,
+                                     Value source) {
+  auto sourceType = source.getType().cast<MemRefType>();
+  unsigned sourceRank = sourceType.getRank();
+  IndexType indexType = builder.getIndexType();
+  SmallVector<Type> sizeStrideTypes(sourceRank, indexType);
+  auto memrefType =
+      MemRefType::get({}, sourceType.getElementType(),
+                      MemRefLayoutAttrInterface{}, sourceType.getMemorySpace());
+  ExtractStridedMetadataOp::build(builder, state, memrefType, indexType,
+                                  sizeStrideTypes, sizeStrideTypes, source);
+}
+
 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 &&