[mlir][Linalg] Introduce linalg.pad_tensor op.

[hanchung]

Summary

linalg.pad_tensor is an operation that pads the source tensor
with given low and high padding config.

Example 1:

mlir
  %pad_value = ... : f32
  %1 = linalg.pad_tensor %0 low[1, 2] high[2, 3] {
  ^bb0(%arg0 : index, %arg1 : index):
    linalg.yield %pad_value : f32
  } : tensor<?x?xf32> to tensor<?x?xf32>

Example 2:

mlir
  %pad_value = ... : f32
  %1 = linalg.pad_tensor %arg0 low[2, %arg1, 3, 3] high[3, 3, %arg1, 2] {
  ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index):
    linalg.yield %pad_value : f32
  } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

52

This backquote isn't closed.

60

You don't make it clear if the source buffer is *modified* by this op.
It seems like it is, because I don't see how to store the padded values otherwise, but that's not great that an operation which is described as "taking a view" is actually modifying the source buffer.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

60

Thanks Mehdi (Hanhan and I discussed this op extensively and the implementation is based on that discussion). Would like to hash out semantics of this.

The op is not modifying the source buffer while taking a view. Instead it is providing a view of the underlying buffer with information about what values to use if you go "outside" of the underlying source buffer", i.e.

• When you read from the view, if you are not in the padding, you get the value from the source buffer. If you are in the padding you get the padded value specified using the region of the operation
• When you write from the view, if you are not in the padding, you write to the source buffer. If you are in the padding the write is ignored. (@Hanhan maybe we should make this semantics explicit)

For the write I agree the semantics is strange (but arguably correct). I dont expect the padded_view to be used to write data. Its only used to read, but the semantics is fairly easy to specify.
This is a way to not always require creating a new buffer to implement padding. Using this op you can fold the padding with its load/store which is a more efficient way to implement padding. You can also fold this with vector.transfer_read/transfer_write to used masked operations and vectorize the computation even with padding.

The advantage here is that the transformation to do tiling/vectorization, etc do not need to actually worry about the padding. They can be implemented as if they are working with the padded buffer. The padding is then implemented using rewrites on load/stores.
Goes without saying this is experimental to some extent, but thats what a lot of manual implementations of ops like conv/pooling, etc. implement

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

60

I have some concerns with what you're describing, because this "view" produces a memref that isn't usable without understanding and accessing this padded_view operation to be able to get the padded value.

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

60

Another aspect is that it does not describe when is the region evaluated: all these ops are evaluating their region when the op is executed, while here it seems that your intent is to fold the region to wherever the memref is used. That can yield weird "effects at a distance" when the region does more than just return a fixed SSA value.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

60

I have some concerns with what you're describing, because this "view" produces a memref that isn't usable without understanding and accessing this padded_view operation to be able to get the padded value.

Explicitly trying to avoid that. This is producing a view (which is memref type) and not a new allocation. Within Linalg itself, you dont need to know where this "memref" came from. Can you give me a more concrete example of why the memref produced here is not usable. From my reading of what this sentence says, the same is true even for a subview. So maybe missing something.

Another aspect is that it does not describe when is the region evaluated: all these ops are evaluating their region when the op is executed, while here it seems that your intent is to fold the region to wherever the memref is used. That can yield weird "effects at a distance" when the region does more than just return a fixed SSA value.

The region is evaluated when you try to access the source memref using indices that are within the padded region. I am not sure I understand what you mean by region not returning a fixed SSA value. It has a single return value. We could have the region describe the conditionals logic described below, but thats seems cumbersome and not really required for the op semantics.

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

60

Explicitly trying to avoid that. This is producing a view (which is memref type) and not a new allocation. Within Linalg itself, you dont need to know where this "memref" came from. Can you give me a more concrete example of why the memref produced here is not usable. From my reading of what this sentence says, the same is true even for a subview. So maybe missing something.

As far as I understand it, a subview only transforms the mapping from the virtual index space into the underlying buffer, this isn't the case here because of the padding.
For example you have this example in the test:

%0 = linalg.padded_view %arg0[1, 2] [2, 3] {
  ^bb0(%arg1 : index, %arg2 : index):
    linalg.yield %pad_value : f32
  } : memref<3x4xf32> to memref<6x9xf32>

What can we do with the %0 memref? Can we use it like other memref?
So for example can you pass it to the runtime print?
We have mlir-cpu-runner test where you should be able to do:

%cast = memref_cast %0 : memref<6x9xf32> to memref<*xf32>
call @print_memref_f32(%unranked_input) : (memref<*xf32>) -> ()

Will this print a 6x9 output? What will it print for the padding?

The region is evaluated when you try to access the source memref using indices that are within the padded region.

As I mentioned before, this seems dangerously fragile to me. At minima that requires ensuring that the region has no side effects.

I am not sure I understand what you mean by region not returning a fixed SSA value. It has a single return value.

I mean that the SSA value alone is not enough, the SSA value somehow carries with it the closure that is represented by the region. This isn't consistent with memref in general I believe.
But we can clarify this point by looking at the example I provide above, this isn't a different point here.

[hanchung]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

60

I don't know how print_memref_f32 work, but I can get a bit what you're saying. Different from other memref, we need a special logic to handle the output of linalg.padded_view. linalg.padded_view not only maps the access indices into the buffer, it also defines non-mapped elements in the buffer.

I feel theoretically it is correct because what print_memref_f32 does is to iterate over all the elements like:

scf.for %iv0 = ... {
  scf.for %iv1 = ... {
    %0 = load padded_view[%iv0, %iv1] : f32
    print %0 : f32
  }
}

And a valid lowering of padded_view is to fold it into if-else to extract the load_value or pad_value. So it will print a 6x9 output.

As Mahesh mentioned, this operation is expected to work within Linalg itself. After applying some transforms, the operation will just fold into either if-else op or alloc + fill + copy. The goal to create a Linalg operation which can represent "pad" semantics in Linalg, work with Linalg transforms, and the op will get killed at some point. I would expect to apply passes (like linalg-to-std) to work with mlir-cpu-runner.

I agree that we can have a more complete definition on the op like Mahesh stated:

When you read from the view, if you are not in the padding, you get the value from the source buffer. If you are in the padding you get the padded value specified using the region of the operation.
When you write from the view, if you are not in the padding, you write to the source buffer. If you are in the padding the write is ignored.

Having a region is easier to extend the pad op to handle different padding cases, like repeat_edge, mirror_edge, etc. But if this is not consistent with memref, I think we can have an explicit operation like padded_scalar_view without taking a region.

[hanchung]

Add more descriptions

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

60

I feel theoretically it is correct because what print_memref_f32 does is to iterate over all the elements like:
...
And a valid lowering of padded_view is to fold it into if-else to extract the load_value or pad_value. So it will print a 6x9 output.

I assume that it is %padded_view here and it is the SSA value returned by linalg.padded_view?

Then it would require *every* consumer of the memref to look for the producer and understand how it need to handle it: the SSA value you're producing is not a valid memref by itself, this is a problem to me.

print_memref_f32 is implemented here: https://github.com/llvm/llvm-project/blob/master/mlir/lib/ExecutionEngine/RunnerUtils.cpp#L39

[bondhugula]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

40

pad -> padded_view

41

Nit: : after example

45

This is where you could really use the custom syntax to improve readability. Consider using keywords like low and high so that it's clear what corresponds to low/high and what corresponds to dimensions. Eg:
linalg.padded_view %0 low [1, 2] high [2, 3] {

45

Unfortunately, this design is flawed as @mehdi_amini hints in his first message -- simply because the resulting memref type isn't carrying the necessary information. This basically means that one is lost if the memref "escapes" (not just interprocedurally but also within a function in various ways), passed along to nested regions with arguments or propagated along in explicit capture style, or in numerous other ways. This is a common pitfall of not having the necessary information encoded in the type. @mravishankar - note that this is NOT a problem with the other view/subview like ops or any other memref creating cast ops that I've seen till date --- because the necessary information to lower say the load/stores is available in the *type*. If other memref defining ops with behavior like these have been introduced where you have to look at the defining operation to see what's happening with accesses (like @mehdi_amini points out), I believe that's equally grave!

51

The first list ...

[mravishankar]

Catching up on the discussion after break here. I agree the design is flawed. The strides are wrong of the result memref. As defined right now that would mean that lowering the memref will need to look at where the memref is coming from which is really wrong. You could return a memref with the strides embedded into the memref type so that it works fine, but I dont think that fully works either. I think Hanhan had a fix to this shortly. Thanks for the feedback!

[hanchung]

Rework and define the op on tensor.

[hanchung]

Update the doc a bit.

[hanchung]

After a long offline discussion, we have a plan to make pad op work with Linalg transforms. I will add two pad op, one is tensor version and another is memref version. I'll start with adding pad_tensor op to Linalg.

I agree that the original semantic introduces issues on memref. The information is not enough. The current plan is to pass an extra destination memref instead of returning one. I will work on it and send out for review in another patch. Thanks for the feedbacks, they are really helpful!

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

45

Added low and high. This is a good idea, thanks!

[nicolasvasilache]

Thanks for pushing on this Hanhan.
The impl. needs to be a little more involved to allow different numbers of low and high padding values.
You'll also need extra builders and accessors to make things easier to manipulate, they can be added on a per need basis.
You can look at Subview/Subtensor and the OffsetSizesAndStridesInterface for similarly looking code.
Refactorings to reuse code are most welcome, if reasonable.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

121

I think you can't just use SameVariadicOperandSize, the following shouldn't work:
.pad_tensor %t low[0, 0] high[%ub0, %ub1].

Indeed you're missing a test for this.

You need to handle your operand_segment_size manually.

See https://github.com/llvm/llvm-project/blob/b1e1bbae0e30c89251940efb0780eee6a1b79ecd/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td#L226
and https://github.com/llvm/llvm-project/blob/118a71565462db41cab1dbb0349200627d6e8524/mlir/lib/Interfaces/ViewLikeInterface.cpp#L161 if you need an example.

127

I'd move the examples after the textual description.

mlir/test/Dialect/Linalg/roundtrip.mlir

21

Please add an asymmetrical test as discussed above.

[hanchung]

Address comments

• Use AttrSizedOperandSegments trait
• Add an asymmetrical test
• Add couple extra builders
• Add operands description
• Move examples after the textual description

[nicolasvasilache]

This looks good.
Can you please add one test per verifier failure to test/Dialect/Linalg/invalid.mlir ?
Once these are in, this is good to go.

Thanks @hanchung !

[hanchung]

Add one more test to invalid.mlir

[hanchung]

In D93704#2512931, @nicolasvasilache wrote:

This looks good.
Can you please add one test per verifier failure to test/Dialect/Linalg/invalid.mlir ?
Once these are in, this is good to go.

Thanks @hanchung !

Added one more test.

The only missing one is for multiple blocks. I think it is not easy to test in invalid.mlir because the parser doesn't parse two regions.

[hanchung]

Add no block test

[hanchung]

In D93704#2513404, @hanchung wrote:

In D93704#2512931, @nicolasvasilache wrote:

This looks good.
Can you please add one test per verifier failure to test/Dialect/Linalg/invalid.mlir ?
Once these are in, this is good to go.

Thanks @hanchung !

Added one more test.

The only missing one is for multiple blocks. I think it is not easy to test in invalid.mlir because the parser doesn't parse two regions.

I was wrong, added.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

class Linalg_Op<string mnemonic, list<OpTrait> traits = []> :
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 def Linalg_InitTensorOp : Linalg_Op<"init_tensor", [NoSideEffect]> {
   let summary = "operation to define a tensor of particular value";
 
   let description = [{
     `linalg.init_tensor` is an operation that materializes a tensor of
     a given shape. The shape could be dynamic or static.
   }];
 
   let arguments =
     (ins Variadic<Index>:$sizes, I64ArrayAttr:$static_sizes);
 
   let results = (outs AnyTensor:$result);
 
   let verifier = [{ return ::verify(*this); }];
 
   let extraClassDeclaration = [{
     static StringRef getStaticSizesAttrName() {
       return "static_sizes";
     }
 
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>(); }
 
     // Infer the shape of the result tensor given the static shapes
     // and element type of the result tensor.
     static Type inferResultType(ArrayRef<int64_t> staticSizes, Type elementType);
 
     // Return true if the size of the tensor is dynamic at `idx`
     bool isDynamicSize(unsigned idx) {
       APInt v = *(static_sizes().getAsValueRange<IntegerAttr>().begin() + idx);
       return ShapedType::isDynamic(v.getSExtValue());
     }
 
     // Assert that the size of the result tensor is static at `idx`
let builders = [
     [{
       build($_builder, $_state, ValueRange{}, staticShape, elementType);
     }]>
   ];
 
   let hasCanonicalizer = 1;
 }
 
+def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
+    [AttrSizedOperandSegments, SingleBlockImplicitTerminator<"YieldOp">]> {
+  let summary = "tensor pad operation";
+  let description = [{
+    `linalg.pad_tensor` is an operation that pads the `source` tensor
+    with given `low` and `high` padding config.
+
+    The PadTensor operation supports the following arguments:
+
+    * source: the "base" tensor on which to pad.
+    * low: A list contains the padding along the start of each
+           dimension, i.e `low`.
+    * high: A list contains the padding along the end of each
+           dimension, i.e. `high`.
+
+    The result tensor dimensions are `low` + `dim` + `high` along that
+    dimension. The number of elements of `low` and `high` must match
+    the rank of the input tensor (which is also the rank of the output
+    tensor). They can be either a constant or a dynamic value.
+
+    The region of the `pad_tensor` operation returns the value to use
+    for the padding. The arguments of the region represent the index
+    of the source being accessed. There should be as many arguments as
+    the rank of the `source` tensor. The value `yield`-ed by the
+    region is used as the value of the view at the given position.
+
+    Example 1:
+
+    ```mlir
+      %pad_value = ... : f32
+      %0 = linalg.pad_tensor %0 low[1, 2] high[2, 3] {
+      ^bb0(%arg0 : index, %arg1 : index):
+        linalg.yield %pad_value : f32
+      } : tensor<?x?xf32> to tensor<?x?xf32>
+    ```
+
+    Example 2:
+
+    ```mlir
+      %pad_value = ... : f32
+      %0 = linalg.pad_tensor %arg0 low[2, %arg1, 3, 3] high[3, 3, %arg1, 2] {
+      ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index):
+          linalg.yield %pad_value : f32
+      } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
+    ```
+
+    Example 3:
+
+    ```mlir
+      %pad_value = ... : f32
+      %0 = linalg.pad_tensor %arg0 low[0, 0] high[%ub0, %ub1] {
+      ^bb0(%arg1: index, %arg2: index):
+        linalg.yield %pad_value : f32
+      } : tensor<2x3xf32> to tensor<?x?xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    AnyTensor:$source,
+    Variadic<Index>:$low,
+    Variadic<Index>:$high,
+    I64ArrayAttr:$static_low,
+    I64ArrayAttr:$static_high);
+
+  let regions = (region AnyRegion:$region);
+
+  let results = (outs AnyTensor:$result);
+
+  let extraClassDeclaration = [{
+    static StringRef getStaticLowAttrName() {
+      return "static_low";
+    }
+
+    static StringRef getStaticHighAttrName() {
+      return "static_high";
+    }
+
+    // Infer the shape of the result tensor given the static shapes
+    // and element type of the result tensor.
+    static RankedTensorType inferResultType(RankedTensorType sourceType,
+                                ArrayRef<int64_t> staticLow,
+                                ArrayRef<int64_t> staticHigh);
+  }];
+
+  let builders = [
+    // Build a PadTensorOp with mixed static and dynamic entries.
+    OpBuilderDAG<(ins "Value":$source, "ArrayRef<int64_t>":$staticLow,
+      "ArrayRef<int64_t>":$staticHigh, "ValueRange":$low, "ValueRange":$high,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
+    // Build a PadTensorOp with all dynamic entries.
+    OpBuilderDAG<(ins "Value":$source, "ValueRange":$low, "ValueRange":$high,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
+  ];
+}
+
 def Linalg_RangeOp :
     Linalg_Op<"range", [NoSideEffect]>,
     Arguments<(ins Index:$min, Index:$max, Index:$step)>,
     Results<(outs Range)> {
   let summary = "Create a `range` type value, used to create `view`s";
   let description = [{
     The `linalg.range` op creates a `!linalg.range` from 3 values of type
     `index` that represent the min, max and step values of the `range`. This

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 
 void InitTensorOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<FoldWithTensorReshapeOp, ReplaceDimOfInitTensorOp,
                  ReplaceStaticShapeDims>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// PadTensorOp
+//===----------------------------------------------------------------------===//
+
+/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
+static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
+  return llvm::to_vector<4>(
+      llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
+        return a.cast<IntegerAttr>().getInt();
+      }));
+}
+
+static LogicalResult verify(PadTensorOp op) {
+  auto sourceType = op.source().getType().cast<RankedTensorType>();
+  auto resultType = op.result().getType().cast<RankedTensorType>();
+  auto expectedType = PadTensorOp::inferResultType(
+      sourceType, extractFromI64ArrayAttr(op.static_low()),
+      extractFromI64ArrayAttr(op.static_high()));
+  if (resultType != expectedType) {
+    return op.emitError("specified type ")
+           << resultType << " does not match the inferred type "
+           << expectedType;
+  }
+
+  auto &region = op.region();
+  if (!llvm::hasSingleElement(region))
+    return op.emitOpError("expected region with 1 block");
+  unsigned rank = resultType.getRank();
+  Block &block = region.front();
+  if (block.getNumArguments() != rank)
+    return op.emitError("expected the block to have ") << rank << " arguments";
+
+  // Note: the number and type of yield values are checked in the YieldOp.
+  for (auto en : llvm::enumerate(block.getArgumentTypes())) {
+    if (!en.value().isIndex())
+      return op.emitOpError("expected block argument ")
+             << (en.index() + 1) << " to be an index";
+  }
+
+  return success();
+}
+
+RankedTensorType PadTensorOp::inferResultType(RankedTensorType sourceType,
+                                              ArrayRef<int64_t> staticLow,
+                                              ArrayRef<int64_t> staticHigh) {
+  unsigned rank = sourceType.getRank();
+  assert(staticLow.size() == rank && "unexpected staticLow size mismatch");
+  assert(staticHigh.size() == rank && "unexpected staticHigh size mismatch");
+
+  SmallVector<int64_t, 4> resultShape;
+  for (auto i : llvm::seq<unsigned>(0, rank)) {
+    if (sourceType.isDynamicDim(i) ||
+        staticLow[i] == ShapedType::kDynamicSize ||
+        staticHigh[i] == ShapedType::kDynamicSize) {
+      resultShape.push_back(ShapedType::kDynamicSize);
+    } else {
+      int64_t size = sourceType.getDimSize(i) + staticLow[i] + staticHigh[i];
+      resultShape.push_back(size);
+    }
+  }
+
+  return RankedTensorType::get(resultShape, sourceType.getElementType());
+}
+
+static ParseResult parsePadTensorOp(OpAsmParser &parser,
+                                    OperationState &result) {
+  OpAsmParser::OperandType baseInfo;
+  SmallVector<OpAsmParser::OperandType, 8> operands;
+  SmallVector<Type, 8> types;
+  if (parser.parseOperand(baseInfo))
+    return failure();
+
+  IndexType indexType = parser.getBuilder().getIndexType();
+  SmallVector<OpAsmParser::OperandType, 4> lowPadding, highPadding;
+  if (parser.parseKeyword("low") ||
+      parseListOfOperandsOrIntegers(parser, result,
+                                    PadTensorOp::getStaticLowAttrName(),
+                                    ShapedType::kDynamicSize, lowPadding))
+    return failure();
+  if (parser.parseKeyword("high") ||
+      parseListOfOperandsOrIntegers(parser, result,
+                                    PadTensorOp::getStaticHighAttrName(),
+                                    ShapedType::kDynamicSize, highPadding))
+    return failure();
+
+  SmallVector<OpAsmParser::OperandType, 8> regionOperands;
+  std::unique_ptr<Region> region = std::make_unique<Region>();
+  SmallVector<Type, 8> operandTypes, regionTypes;
+  if (parser.parseRegion(*region, regionOperands, regionTypes))
+    return failure();
+  result.addRegion(std::move(region));
+
+  Type srcType, dstType;
+  if (parser.parseColonType(srcType) || parser.parseKeywordType("to", dstType))
+    return failure();
+
+  if (parser.addTypeToList(dstType, result.types))
+    return failure();
+
+  SmallVector<int, 4> segmentSizesFinal = {1}; // source tensor
+  segmentSizesFinal.append({static_cast<int>(lowPadding.size()),
+                            static_cast<int>(highPadding.size())});
+  result.addAttribute(
+      OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr(),
+      parser.getBuilder().getI32VectorAttr(segmentSizesFinal));
+  return failure(
+      parser.parseOptionalAttrDict(result.attributes) ||
+      parser.resolveOperand(baseInfo, srcType, result.operands) ||
+      parser.resolveOperands(lowPadding, indexType, result.operands) ||
+      parser.resolveOperands(highPadding, indexType, result.operands));
+}
+
+static void print(OpAsmPrinter &p, PadTensorOp op) {
+  p << op->getName().getStringRef() << ' ';
+  p << op.source();
+  p << " low";
+  printListOfOperandsOrIntegers(p, op.low(), op.static_low(),
+                                ShapedType::isDynamic);
+  p << " high";
+  printListOfOperandsOrIntegers(p, op.high(), op.static_high(),
+                                ShapedType::isDynamic);
+  p.printRegion(op.region());
+  p << " : " << op.source().getType() << " to " << op.getType();
+}
+
+void PadTensorOp::build(OpBuilder &b, OperationState &result, Value source,
+                        ArrayRef<int64_t> staticLow,
+                        ArrayRef<int64_t> staticHigh, ValueRange low,
+                        ValueRange high, 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));
+  result.addAttributes(attrs);
+}
+
+void PadTensorOp::build(OpBuilder &b, OperationState &result, Value source,
+                        ValueRange low, ValueRange high,
+                        ArrayRef<NamedAttribute> attrs) {
+  auto sourceType = source.getType().cast<RankedTensorType>();
+  unsigned rank = sourceType.getRank();
+  SmallVector<int64_t, 4> staticVector(ShapedType::kDynamicSize, rank);
+  build(b, result, source, staticVector, staticVector, low, high, attrs);
+}
+
+//===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 /// Collapse reassociation maps that are used in pair of reshape ops where one
 /// is a producer and other is the consumer. Only valid to use this method when
 /// both the producer and consumer are collapsing dimensions or both are
 /// expanding dimensions.
 ///

 static LogicalResult verify(linalg::YieldOp op) {
   auto *parentOp = op->getParentOp();
   if (parentOp->getNumRegions() != 1 || parentOp->getRegion(0).empty())
     return op.emitOpError("expected single non-empty parent region");
 
   if (auto linalgOp = dyn_cast<LinalgOp>(parentOp))
     return verifyYield(op, cast<LinalgOp>(parentOp));
 
+  if (auto padTensorOp = dyn_cast<linalg::PadTensorOp>(parentOp)) {
+    return success(
+        op.getNumOperands() == 1 &&
+        op.getOperand(0).getType() ==
+            padTensorOp.getType().cast<ShapedType>().getElementType());
+  }
+
   return op.emitOpError("expected parent op with LinalgOp interface");
 }
 
 /////// Operations corresponding to library calls defined with Tablegen ////////
 
 void FillOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {

mlir/test/Dialect/Linalg/invalid.mlir

 {
   // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
   %0 = linalg.reshape %arg0
          [affine_map<(d0, d1, d2) -> (d0)>,
           affine_map<(d0, d1, d2) -> (d1, d2)>] :
        memref<?x4x5xf32> into memref<?x?xf32>
   return %0 : memref<?x?xf32>
 }
+
+// -----
+
+func @pad_result_type(%arg0: tensor<?x2x3x4xi32>, %arg1: index, %arg2: i32) -> tensor<?x?x?x8xf32> {
+  // expected-error @+1 {{specified type 'tensor<?x?x?x8xf32>' does not match the inferred type 'tensor<?x?x?x9xi32>}}
+  %0 = linalg.pad_tensor %arg0 low[1, %arg1, 2, 2] high[1, 2, %arg1, 3] {
+  ^bb0(%arg3: index, %arg4: index):  // no predecessors
+    linalg.yield %arg2 : i32
+  } : tensor<?x2x3x4xi32> to tensor<?x?x?x8xf32>
+  return %0 : tensor<?x?x?x8xf32>
+}
+
+// -----
+
+func @pad_number_of_block_args(%arg0: tensor<?x4xi32>, %arg1: i32) -> tensor<?x9xi32> {
+  // expected-error @+1 {{expected the block to have 2 arguments}}
+  %0 = linalg.pad_tensor %arg0 low[1, 2] high[2, 3] {
+  ^bb0(%arg2: index, %arg3: index, %arg4: index):  // no predecessors
+    linalg.yield %arg1 : i32
+  } : tensor<?x4xi32> to tensor<?x9xi32>
+  return %0 : tensor<?x9xi32>
+}
+
+// -----
+
+func @pad_no_block(%arg0: tensor<?x4xi32>, %arg1: i32) -> tensor<?x9xi32> {
+  // expected-error @+1 {{expected region with 1 block}}
+  %0 = linalg.pad_tensor %arg0 low[1, 2] high[2, 3] {
+  } : tensor<?x4xi32> to tensor<?x9xi32>
+  return %0 : tensor<?x9xi32>
+}
+
+// -----
+
+func @pad_block_args(%arg0: tensor<?x4xi32>, %arg1: i32) -> tensor<?x9xi32> {
+  // expected-error @+1 {{op expected block argument 1 to be an index}}
+  %0 = linalg.pad_tensor %arg0 low[1, 2] high[2, 3] {
+  ^bb0(%arg2: i32, %arg3: i32):  // no predecessors
+    linalg.yield %arg1 : i32
+  } : tensor<?x4xi32> to tensor<?x9xi32>
+  return %0 : tensor<?x9xi32>
+}

mlir/test/Dialect/Linalg/roundtrip.mlir

 // RUN: mlir-opt -split-input-file %s | FileCheck %s
 
 // TODO: Re-enable LLVM lowering test after IndexedGenericOp is lowered.
 //
 // Test that we can lower all the way to LLVM without crashing, don't check results here.
 // DISABLED: mlir-opt %s --convert-linalg-to-llvm -o=/dev/null 2>&1
 
+func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
+                  %pad_value: f32) -> tensor<6x?x?x?xf32> {
+  %0 = linalg.pad_tensor %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
+    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
+      linalg.yield %pad_value : f32
+    } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
+  return %0 : tensor<6x?x?x?xf32>
+}
+// CHECK-LABEL: func @pad_dynamic
+//  CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]
+//  CHECK-SAME: %[[LOW:[a-zA-Z0-9_]*]]
+//  CHECK-SAME: %[[HIGH:[a-zA-Z0-9_]*]]
+//       CHECK:   linalg.pad_tensor %[[ARG0]]
+//  CHECK-SAME:     low[2, %[[LOW]], 3, 3]
+//  CHECK-SAME:     high[3, 3, %[[HIGH]], 2]
+//       CHECK:    : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
+
+// -----
+
+func @pad_static(%arg0: tensor<3x4xf32>, %pad_value: f32) -> tensor<6x9xf32> {
+  %0 = linalg.pad_tensor %arg0 low[1, 2] high[2, 3] {
+    ^bb0(%arg1 : index, %arg2 : index):
+      linalg.yield %pad_value : f32
+    } : tensor<3x4xf32> to tensor<6x9xf32>
+  return %0 : tensor<6x9xf32>
+}
+// CHECK-LABEL: func @pad_static
+//  CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]
+//       CHECK:   linalg.pad_tensor %[[ARG0]] low[1, 2] high[2, 3]
+//       CHECK:    : tensor<3x4xf32> to tensor<6x9xf32>
+
+// -----
+
+func @pad_asymmetrical(%arg0: tensor<2x3xf32>, %ub0: index, %ub1: index,
+                       %pad_value: f32) -> tensor<?x?xf32> {
+  %0 = linalg.pad_tensor %arg0 low[0, 0] high[%ub0, %ub1] {
+    ^bb0(%arg1: index, %arg2: index):
+      linalg.yield %pad_value : f32
+    } : tensor<2x3xf32> to tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @pad_asymmetrical
+//  CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]
+//  CHECK-SAME: %[[UB0:[a-zA-Z0-9_]*]]
+//  CHECK-SAME: %[[UB1:[a-zA-Z0-9_]*]]
+//       CHECK:   linalg.pad_tensor %[[ARG0]]
+//  CHECK-SAME:     low[0, 0]
+//  CHECK-SAME:     high[%[[UB0]], %[[UB1]]]
+//       CHECK:    : tensor<2x3xf32> to tensor<?x?xf32>
+
+// -----
+
 func @range(%arg0: index, %arg1: index, %arg2: index) {
   %0 = linalg.range %arg0:%arg1:%arg2 : !linalg.range
   return
 }
 // CHECK-LABEL: func @range(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 //  CHECK-NEXT:  linalg.range %{{.*}} : %{{.*}} : %{{.*}} : !linalg.range
 
 // -----