[mlir][tensor] Implement ReifyRankedShapedTypeOpInterface on PadOp

[pifon2a]

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

1834

nit: cache getLoc()?

[springerm]

Just noticed that this is already implemented as an external model. (mlir::tensor::registerInferTypeOpInterfaceExternalModels)

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

def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // PadOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_PadOp : Tensor_Op<"pad", [AttrSizedOperandSegments, NoSideEffect,
+    DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
     SingleBlockImplicitTerminator<"mlir::tensor::YieldOp">]> {
   let summary = "tensor pad operation";
   let description = [{
     `tensor.pad` is an operation that pads the `source` tensor
     with given `low` and `high` padding config.
 
     The PadOp operation supports the following arguments:
 

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

if (sourceType.isDynamicDim(i) ||
              "mismatch between inferred shape and result shape");
       inferredShape.push_back(size);
     }
   }
 
   return RankedTensorType::get(inferredShape, sourceType.getElementType());
 }
 
+LogicalResult
+PadOp::reifyResultShapes(OpBuilder &builder,
+                         ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+  RankedTensorType resultType = getResultType();
+  reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
+
+  // Given an OpFoldResult, return an index-typed value.
+  auto getIdxValue = [&](OpFoldResult ofr) {
+    if (auto val = ofr.dyn_cast<Value>())
+      return val;
+    return builder
+        .create<arith::ConstantIndexOp>(getLoc(), *getConstantIntValue(ofr))
+        .getResult();
+  };
+
+  for (auto dim : llvm::seq<int64_t>(0, resultType.getRank())) {
+    if (!resultType.isDynamicDim(dim)) {
+      // Static result dim. No need to compute anything.
+      reifiedReturnShapes[0][dim] = builder.create<arith::ConstantIndexOp>(
+          getLoc(), resultType.getDimSize(dim));
+      continue;
+    }
+
+    // Compute src size + low pad + high pad.
+    Value srcDimSize =
+        builder.createOrFold<tensor::DimOp>(getLoc(), source(), dim);
+    Value lowPadSize = getIdxValue(getMixedLowPad()[dim]);
+    Value highPadSize = getIdxValue(getMixedHighPad()[dim]);
+    reifiedReturnShapes[0][dim] = builder.createOrFold<arith::AddIOp>(
+        getLoc(), srcDimSize,
+        builder.createOrFold<arith::AddIOp>(getLoc(), lowPadSize, highPadSize));
+  }
+  return success();
+}
+
 void PadOp::build(OpBuilder &b, OperationState &result, Value source,
                   ArrayRef<int64_t> staticLow, ArrayRef<int64_t> staticHigh,
                   ValueRange low, ValueRange high, bool nofold,
                   ArrayRef<NamedAttribute> attrs) {
   auto sourceType = source.getType().cast<RankedTensorType>();
   auto resultType = inferResultType(sourceType, staticLow, staticHigh);
   build(b, result, resultType, source, low, high, b.getI64ArrayAttr(staticLow),
         b.getI64ArrayAttr(staticHigh), nofold ? b.getUnitAttr() : UnitAttr());