[mlir][tensor] Implement ReifyRankedShapedTypeOpInterface on GenerateOp

[springerm]

Summary

Depends On D121519

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

 }
 
 //===----------------------------------------------------------------------===//
 // GenerateOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_GenerateOp : Tensor_Op<"generate",
     [RecursiveSideEffects,
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
      SingleBlockImplicitTerminator<"mlir::tensor::YieldOp">]> {
   string summary = "Creates a dynamically sized tensor from elements";
   string description = [{
     This operation creates a dynamically sized tensor with elements of any type.
     It expects one index operand per dynamic extent of the result tensor.
 
     The body region defines the tensor's elements. It takes index operands as
     its region arguments that span the index space. The element at the given

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

if (auto splatDest = dest.dyn_cast<SplatElementsAttr>())
         return dest;
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // GenerateOp
 //===----------------------------------------------------------------------===//
 
+LogicalResult GenerateOp::reifyResultShapes(
+    OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+  reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
+  int idx = 0;
+  for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
+    if (getType().isDynamicDim(dim)) {
+      reifiedReturnShapes[0][dim] = getOperand(idx++);
+    } else {
+      reifiedReturnShapes[0][dim] = builder.create<arith::ConstantIndexOp>(
+          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.
   RankedTensorType resultTy = getType().cast<RankedTensorType>();
   if (getNumOperands() != resultTy.getNumDynamicDims())
     return emitError("must have as many index operands as dynamic extents "
                      "in the result type");