[mlir][Tensor] Implement `reifyReturnTypeShapesPerResultDim` for `tensor.insert_slice`.

[cathyzhyi]

LGTM

[mravishankar]

Rebase

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

 #define MLIR_DIALECT_TENSOR_IR_TENSOR_H_
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 //===----------------------------------------------------------------------===//
 // Tensor Dialect Helpers
 //===----------------------------------------------------------------------===//
 
 namespace mlir {

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

 //===- TensorOps.td - Tensor 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 TENSOR_OPS
 #define TENSOR_OPS
 
 include "mlir/Dialect/Tensor/IR/TensorBase.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 class Tensor_Op<string mnemonic, list<OpTrait> traits = []>
     : Op<Tensor_Dialect, mnemonic, traits> {
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
let arguments = (ins AnyTensor:$source,
                        Index:$index);
   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)>
+    OpBuilder<(ins "Value":$source, "int64_t":$index)>
   ];
 
   let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     Optional<int64_t> getConstantIndex();
   }];
 
   let hasCanonicalizer = 1;

 
 //===----------------------------------------------------------------------===//
 // InsertSliceOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_InsertSliceOp : BaseOpWithOffsetSizesAndStrides<
     Tensor_Dialect, "insert_slice",
     [NoSideEffect, AttrSizedOperandSegments, OffsetSizeAndStrideOpInterface,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+         ["reifyReturnTypeShapesPerResultDim"]>,
      TypesMatchWith<"expected result type to match dest type",
                     "dest", "result", "$_self">]> {
   let summary = "insert_slice operation";
   let description = [{
     The "insert_slice" operation insert a tensor `source` into another
     tensor `dest` as specified by the operation's offsets, sizes and strides
     arguments.
 

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRTensor
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRCastInterfaces
   MLIRDialectUtils
   MLIRIR
+  MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRStandard
   MLIRViewLikeInterface
   )

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 
 void DimOp::build(OpBuilder &builder, OperationState &result, Value source,
                   int64_t index) {
   auto loc = result.location;
   Value indexValue = builder.create<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 = index().getDefiningOp<ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 
 static LogicalResult verify(DimOp op) {
   // Assume unknown index to be in range.

 OpFoldResult InsertSliceOp::fold(ArrayRef<Attribute>) {
   if (getSourceType().hasStaticShape() && getType().hasStaticShape() &&
       getSourceType() == getType() &&
       succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
     return this->source();
   return OpFoldResult();
 }
 
+LogicalResult InsertSliceOp::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &builder,
+    SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
+  for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
+    reifiedReturnShapes[0][dim] =
+        builder.createOrFold<tensor::DimOp>(getLoc(), dest(), dim);
+  }
+  return success();
+}
+
 namespace {
 /// Pattern to rewrite a insert_slice op with constant arguments.
 class InsertSliceOpConstantArgumentFolder final
     : public OpRewritePattern<InsertSliceOp> {
 public:
   using OpRewritePattern<InsertSliceOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp,

mlir/test/Dialect/Tensor/resolve-shaped-type-result-dims.mlir

+// RUN: mlir-opt -resolve-shaped-type-result-dims -split-input-file %s | FileCheck %s
+
+func @insert_slice(
+    %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>,
+    %arg2 : index, %arg3 : index, %arg4 : index) -> (index, index, index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+  %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
+  %0 = tensor.insert_slice %arg0 into %arg1[%arg2, %arg3, %arg4] [%d0, %d1, %d2] [1, 1, 1] : tensor<?x?x?xf32> into tensor<?x?x?xf32>
+  %1 = tensor.dim %0, %c0 : tensor<?x?x?xf32>
+  %2 = tensor.dim %0, %c1 : tensor<?x?x?xf32>
+  %3 = tensor.dim %0, %c2 : tensor<?x?x?xf32>
+  return %1, %2, %3 : index, index, index
+}
+// CHECK-LABEL: func @insert_slice(
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//   CHECK-DAG:   %[[C2:.+]] = constant 2 : index
+//   CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[ARG1]], %[[C0]]
+//   CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[ARG1]], %[[C1]]
+//   CHECK-DAG:   %[[D2:.+]] = tensor.dim %[[ARG1]], %[[C2]]
+//       CHECK:   return %[[D0]], %[[D1]], %[[D2]]

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

td_library(
     srcs = [
         "include/mlir/Dialect/Tensor/IR/TensorBase.td",
         "include/mlir/Dialect/Tensor/IR/TensorOps.td",
     ],
     includes = ["include"],
     deps = [
         ":CastInterfacesTdFiles",
         ":ControlFlowInterfacesTdFiles",
+        ":InferTypeOpInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
         ":ViewLikeInterfaceTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "TensorOpsIncGen",
cc_library(
     ) + ["include/mlir/Transforms/InliningUtils.h"],
     hdrs = ["include/mlir/Dialect/Tensor/IR/Tensor.h"],
     includes = ["include"],
     deps = [
         ":CastOpInterfaces",
         ":ControlFlowInterfaces",
         ":DialectUtils",
         ":IR",
+        ":InferTypeOpInterface",
         ":SideEffectInterfaces",
         ":StandardOps",
         ":Support",
         ":TensorOpsIncGen",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )