Remove zero-extent inputs from concat

diff --git a/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp b/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp
index 79acf47..774da5f 100644
--- a/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp
+++ b/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp

@@ -1347,6 +1347,41 @@
   }
 };
 
+struct ZeroConcat final : OpRewritePattern<mlir::stablehlo::ConcatenateOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::ConcatenateOp op,
+                                PatternRewriter &rewriter) const override {
+    auto type = dyn_cast<RankedTensorType>(op.getType());
+    if (!type || !type.hasStaticShape())
+      return failure();
+
+    uint64_t axis = op.getDimension();
+    OperandRange origInputs = op.getInputs();
+    SmallVector<Value> nonzeroInputs;
+    for (auto input : origInputs) {
+      auto type = dyn_cast<RankedTensorType>(input.getType());
+      ArrayRef<int64_t> shape = type.getShape();
+      if (axis > shape.size()) {
+        return failure();
+      }
+
+      if (shape[axis] != 0) {
+        nonzeroInputs.push_back(input);
+      }
+    }
+
+    if (nonzeroInputs.size() == origInputs.size()) {
+      return failure();
+    }
+
+    Value new_op = rewriter.create<mlir::stablehlo::ConcatenateOp>(
+        op.getLoc(), nonzeroInputs, /*dimension=*/axis);
+    rewriter.replaceOp(op, new_op);
+    return success();
+  }
+};
+
 // Similar to DotIsMul, this finds the case where a dot general
 // can be represented using a mul operation. This includes possibly making
 // an implicit cast explicit prior the mul.
@@ -1795,7 +1830,7 @@
 
     // Additional canonicalizers that simplify to computationally
     // less-complex operations.
-    patterns.insert<DotToMul>(context);
+    patterns.insert<DotToMul, ZeroConcat>(context);
 
     // Unary elementwise op.
     patterns.insert<

diff --git a/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/test/stablehlo_to_stablehlo.mlir b/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/test/stablehlo_to_stablehlo.mlir
index 544a584..debce84 100644
--- a/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/test/stablehlo_to_stablehlo.mlir
+++ b/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/test/stablehlo_to_stablehlo.mlir

@@ -462,3 +462,15 @@
 
 // CHECK-LABEL: @broadcast_iota_sort_slice_incorrect_dims
 // CHECK-NOT: chlo.top_k
+
+// -----
+
+// CHECK-LABEL: @concat_remove_zero_extents
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9$._-]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9$._-]+]]
+func.func @concat_remove_zero_extents(%arg0: tensor<2x3xi32>, %arg1 : tensor<2x3xi32>, %arg2 : tensor<2x0xi32>) -> tensor<2x6xi32> {
+  %0 = stablehlo.concatenate %arg0, %arg1, %arg2, dim = 1 : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<2x0xi32>) -> tensor<2x6xi32>
+  // CHECK: [[R0:%.+]] = stablehlo.concatenate  %[[ARG0]], %[[ARG1]], dim = 1 : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x6xi32>
+  return %0 : tensor<2x6xi32>
+}
+