[Mlir-commits] [mlir] fdb41a2 - [mlir][tensor] Implement ReifyRankedShapedTypeOpInterface on GenerateOp

[Matthias Springer]

Author: Matthias Springer
Date: 2022-03-16T18:59:27+09:00
New Revision: fdb41a22ac33edc29703ec9545a451c07bfcfdf8

URL: https://github.com/llvm/llvm-project/commit/fdb41a22ac33edc29703ec9545a451c07bfcfdf8
DIFF: https://github.com/llvm/llvm-project/commit/fdb41a22ac33edc29703ec9545a451c07bfcfdf8.diff

LOG: [mlir][tensor] Implement ReifyRankedShapedTypeOpInterface on GenerateOp

Differential Revision: https://reviews.llvm.org/D121520

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
    mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index c445061c160bd..2021231f28f24 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -372,6 +372,7 @@ def Tensor_FromElementsOp : Tensor_Op<"from_elements", [
 
 def Tensor_GenerateOp : Tensor_Op<"generate",
     [RecursiveSideEffects,
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
      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
index 190c4b1b6bc95..2789fd3b2cd9a 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -525,6 +525,21 @@ OpFoldResult InsertOp::fold(ArrayRef<Attribute> operands) {
 // 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.