Fix failed reify for ReverseOp (#2109)

The `ReverseOp` did not implement a `reifyReturnTypeShapes`. This
resulted in a failed assertion in the `getEmptyTensorFor` helper
function that gets used during the `stablehlo-legalize-to-linalg` pass
when the op has dynamic input/output. This patch implements this
function as well as adds a test for the newly supported case (input can
have dynamic dims as long as they don't correspond to the `dims`
attribute of the op)

Various other quality of life changes, such as having match failure
notifications have been added.

Author: srcarroll

stablehlo/conversions/linalg/tests/miscellaneous.mlir

@@ -1072,6 +1072,21 @@ func.func @reverse(%input: tensor<2x3xf32>) -> tensor<2x3xf32> {
 
 // -----
 
+// CHECK-DAG: #[[OPERAND_MAP:.*]] = affine_map<(d0, d1) -> (d0, -d1 + 2)>
+// CHECK-DAG: #[[RESULT_MAP:.*]]  = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: func @reverse_dynamic
+func.func @reverse_dynamic(%input: tensor<?x3xf32>) -> tensor<?x3xf32> {
+  %result = "stablehlo.reverse"(%input) {
+    dimensions = array<i64: 1>, someattr
+  } : (tensor<?x3xf32>) -> tensor<?x3xf32>
+  func.return %result : tensor<?x3xf32>
+}
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[OPERAND_MAP]], #[[RESULT_MAP]]]
+// CHECK-SAME: {someattr}
+
+// -----
+
 // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d3)>
 // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
 // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>

stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.cpp

@@ -77,7 +77,9 @@ Value getEmptyTensorFor(OpBuilder &b, Location loc, ShapedType resultType,
     // Ask the op for its output shape.
     auto shapeSource = cast<InferShapedTypeOpInterface>(op);
     SmallVector<Value, 1> reifiedShapes;
-    (void)shapeSource.reifyReturnTypeShapes(b, operands, reifiedShapes);
+    assert(succeeded(
+               shapeSource.reifyReturnTypeShapes(b, operands, reifiedShapes)) &&
+           "could not reify");
     assert(reifiedShapes.size() == 1 && "Expected one reified result");
     // Construct sizes for the required dimensions.
     for (const auto &en : llvm::enumerate(resultType.getShape())) {

stablehlo/conversions/linalg/transforms/StablehloLegalizeToLinalg.cpp

@@ -354,18 +354,20 @@ struct DataMovementOpConverter : OpConversionPattern<OpTy> {
   LogicalResult matchAndRewrite(
       OpTy op, typename OpTy::Adaptor adaptor,
       ConversionPatternRewriter &rewriter) const final {
-    if (failed(verifyHloOpBufferOrTensorSemantics(op))) return failure();
+    if (failed(verifyHloOpBufferOrTensorSemantics(op)))
+      return rewriter.notifyMatchFailure(
+          op, "failed to verify hlo buffer or tensor semantics");
 
     ShapedType resultType = getHloOpResultType(op);
     resultType =
         this->getTypeConverter()->template convertType<ShapedType>(resultType);
-    if (!resultType) {
+    if (!resultType)
       return rewriter.notifyMatchFailure(op, "type conversion failed");
-    }
 
     SmallVector<AffineMap, 2> indexingMaps =
         Derived::getIndexingMaps(op, &rewriter);
-    if (indexingMaps.empty()) return failure();
+    if (indexingMaps.empty())
+      return rewriter.notifyMatchFailure(op, "could not derive indexing maps");
 
     int64_t nloops = resultType.getRank();
     Location loc = op.getLoc();

stablehlo/dialect/StablehloOps.cpp

@@ -1605,6 +1605,13 @@ LogicalResult ReverseOp::inferReturnTypeComponents(
                              inferredReturnShapes);
 }
 
+LogicalResult ReverseOp::reifyReturnTypeShapes(
+    OpBuilder& builder, ValueRange operands,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
+  return ::mlir::hlo::deriveShapeFromOperand(
+      &builder, getOperation(), operands.front(), &reifiedReturnShapes);
+}
+
 //===----------------------------------------------------------------------===//
 // RngBitGeneratorOp
 //===----------------------------------------------------------------------===//

stablehlo/dialect/StablehloOps.td

@@ -2730,7 +2730,7 @@ def StableHLO_SortOp : StableHLO_Op<"sort",
   let hasVerifier = 1;
 }
 
-def StableHLO_ReverseOp: StableHLO_Op<"reverse",
+def StableHLO_ReverseOp: StableHLO_ShapedInterfaceOp<"reverse",
       [Pure, HLO_CompatibleOperandsAndResultType /*reverse_c1*/]> {
   let summary = "Reverse operation";
   let description = [{

stablehlo/tests/infer_stablehlo.mlir

@@ -1826,3 +1826,16 @@ func.func @select(%pred : tensor<i1>,
   %1 = "hlo_test_infer.get_return_types"(%0) : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xindex>
   func.return %1 : tensor<?x?x?x?xindex>
 }
+
+// -----
+
+// CHECK-LABEL: @reverse
+// CHECK-SAME:   %[[A:.*]]: tensor<?x?x?x?xf32>
+func.func @reverse(%a : tensor<?x?x?x?xf32>) -> tensor<4xindex> {
+  %0 = "stablehlo.reverse"(%a) {
+    dimensions = array<i64: 1, 3>, someattr
+  } : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[A]] : tensor<?x?x?x?xf32> -> tensor<4xindex>
+  %1 = "hlo_test_infer.reify_return_type_shapes"(%0) : (tensor<?x?x?x?xf32>) -> tensor<4xindex>
+  func.return %1 : tensor<4xindex>
+}