commit 440b8bbac20cd7bc6ae8546e67ff8284d6c8b523
author Sean Silva <silvasean@google.com> Thu Sep 09 20:50:38 2021 -0700
committer GitHub <noreply@github.com> Thu Sep 09 20:50:38 2021 -0700
tree ee00f4ecca73b569ad160a2b4f354f831452ebae
parent a7d7ac2b0bcaa2fcf2b706f084ea8650e65da409

Correctly handle 1:1 conversions for calls in FunctionSignatureExpansionPass (#6997)

The code was copypasted from upstream DecomposeCallGraphTypes, which is
unused and had a latent bug for materializing 1:1 conversions. It turns
out the code can be simplified.

iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp

diff --git a/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp b/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp
index edd0814..7b5b841 100644
--- a/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp
+++ b/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp
@@ -65,16 +65,9 @@
     auto decomposedValues = llvm::to_vector<6>(
         llvm::map_range(expandedResultIndices[i],
                         [&](unsigned i) { return newCallOp.getResult(i); }));
-    if (decomposedValues.empty()) {
-      // No replacement is required.
-      replacedValues.push_back(nullptr);
-    } else if (decomposedValues.size() == 1) {
-      replacedValues.push_back(decomposedValues.front());
-    } else {
-      Value materialized = typeExpander.castToSource(loc, callOp.getType(i),
-                                                     decomposedValues, builder);
-      replacedValues.push_back(materialized);
-    }
+    Value materialized = typeExpander.castToSource(loc, callOp.getType(i),
+                                                   decomposedValues, builder);
+    replacedValues.push_back(materialized);
   }
   callOp.replaceAllUsesWith(replacedValues);
   return success();

iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir

diff --git a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir
index 9ad16d5..5e29f4b 100644
--- a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir
+++ b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir
@@ -43,3 +43,18 @@
   %0 = std.call @calls(%arg0) : (!shapex.ranked_shape<[?,?]>) -> !shapex.ranked_shape<[?,?]>
   return %0 : !shapex.ranked_shape<[?,?]>
 }
+
+// -----
+// CHECK-LABEL:   func @oneUnknownDimension(
+// CHECK-SAME:                              %[[ARG:.*]]: index) -> index {
+// CHECK:           %[[ARG_RS:.*]] = shapex.make_ranked_shape %[[ARG]] : (index) -> !shapex.ranked_shape<[?]>
+// CHECK:           %[[ARG_DIM0:.*]] = shapex.ranked_dim %[[ARG_RS]][0] : !shapex.ranked_shape<[?]> -> index
+// CHECK:           %[[CALL:.*]] = call @oneUnknownDimension(%[[ARG_DIM0]]) : (index) -> index
+// CHECK:           %[[CALL_SHAPE:.*]] = shapex.make_ranked_shape %[[CALL]] : (index) -> !shapex.ranked_shape<[?]>
+// CHECK:           %[[CALL_DIM0:.*]] = shapex.ranked_dim %[[CALL_SHAPE]][0] : !shapex.ranked_shape<[?]> -> index
+// CHECK:           return %[[CALL_DIM0]] : index
+
+func @oneUnknownDimension(%arg0 :!shapex.ranked_shape<[?]>) -> !shapex.ranked_shape<[?]> {
+  %0 = std.call @oneUnknownDimension(%arg0) : (!shapex.ranked_shape<[?]>) -> !shapex.ranked_shape<[?]>
+  return %0 : !shapex.ranked_shape<[?]>
+}