diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td --- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td @@ -372,6 +372,7 @@ def Tensor_GenerateOp : Tensor_Op<"generate", [RecursiveSideEffects, + DeclareOpInterfaceMethods, SingleBlockImplicitTerminator<"mlir::tensor::YieldOp">]> { string summary = "Creates a dynamically sized tensor from elements"; string description = [{ diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp --- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -525,6 +525,21 @@ // GenerateOp //===----------------------------------------------------------------------===// +LogicalResult GenerateOp::reifyResultShapes( + OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { + reifiedReturnShapes.resize(1, SmallVector(getType().getRank())); + int idx = 0; + for (auto dim : llvm::seq(0, getType().getRank())) { + if (getType().isDynamicDim(dim)) { + reifiedReturnShapes[0][dim] = getOperand(idx++); + } else { + reifiedReturnShapes[0][dim] = builder.create( + getLoc(), getType().getDimSize(dim)); + } + } + return success(); +} + LogicalResult GenerateOp::verify() { // Ensure that the tensor type has as many dynamic dimensions as are specified // by the operands.