commit 131f2e27184e7824a5cb85e7b990e415675e9ab8
author Han-Chung Wang <hanchung@google.com> Tue Apr 20 01:03:02 2021 +0800
committer GitHub <noreply@github.com> Tue Apr 20 01:03:02 2021 +0800
tree 8c7c682e4425b164f58b351903fca9d70c37a979
parent 840f4c351295f59760c60fea00745b589ac9ff92

Create ops with static shapes in PadTensorOpConversion if they're static (#5516)

Currently, the mhlo.pad will be lowered to linalg.pad_tensor and then
lowered to `linalg.init_tensor + linalg.fill + subtensor_insert`. The
init_tensor op will produce a dynamic shape even if the shape is static.
This leads a `tensor.cast` op added and relies on further patterns to
fix them. This is not needed to static shape and will hit some issues
related to https://github.com/google/iree/issues/5385.

iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp

diff --git a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp
index cfbc1ae..c7a1327 100644
--- a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp
+++ b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp
@@ -87,7 +87,7 @@
 
     // TODO(ravishankarm): Use shape inference interface to get this.
     SmallVector<OpFoldResult> sourceShape;
-    SmallVector<Value> outputShape;
+    SmallVector<OpFoldResult> outputShape;
     for (int64_t dim : llvm::seq<int64_t>(0, rank)) {
       SmallVector<Value> mapValues;
       Value sourceDim = rewriter.createOrFold<memref::DimOp>(loc, source, dim);
@@ -106,21 +106,21 @@
       };
       expr = addValueOrAttr(expr, lowPad[dim]);
       expr = addValueOrAttr(expr, highPad[dim]);
-      outputShape.push_back(linalg::applyMapToValues(
-          rewriter, loc, AffineMap::get(1, numSymbols, expr), mapValues)[0]);
+      Value v = linalg::applyMapToValues(
+          rewriter, loc, AffineMap::get(1, numSymbols, expr), mapValues)[0];
+      if (auto cst = v.getDefiningOp<ConstantOp>()) {
+        outputShape.push_back(cst.value());
+      } else {
+        outputShape.push_back(v);
+      }
     }
     Value initTensor = rewriter.create<linalg::InitTensorOp>(
         loc, outputShape, sourceType.getElementType());
     Value fill =
         rewriter.create<linalg::FillOp>(loc, initTensor, yieldVal).getResult(0);
     SmallVector<OpFoldResult> strides(rank, rewriter.getI64IntegerAttr(1));
-    Value replacement = rewriter.create<SubTensorInsertOp>(
-        loc, source, fill, lowPad, sourceShape, strides);
-    if (padTensorOp.getResultType() != replacement.getType()) {
-      replacement = rewriter.create<tensor::CastOp>(
-          loc, padTensorOp.getResultType(), replacement);
-    }
-    rewriter.replaceOp(padTensorOp, replacement);
+    rewriter.replaceOpWithNewOp<SubTensorInsertOp>(
+        padTensorOp, source, fill, lowPad, sourceShape, strides);
     return success();
   }
 };