[MLIR][Shape] Lower `shape.any`

[frgossen]

Summary

Lower shape.any to its first operand.

[herhut]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

33

Maybe add a comment here that this is a valid interpretation of the any operation, as this is surprising.

[frgossen]

Add comment

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

 using namespace mlir::shape;
 
 namespace {
 
 /// Generated conversion patterns.
 #include "ShapeToStandardPatterns.inc"
 
 /// Conversion patterns.
+class AnyOpConversion : public OpConversionPattern<AnyOp> {
+public:
+  using OpConversionPattern<AnyOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(AnyOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    AnyOp::Adaptor transformed(operands);
+
+    // Replace `any` with its first operand.
+    // Any operand would be a valid substitution.
+    rewriter.replaceOp(op, {transformed.inputs().front()});
+    return success();
+  }
+};
+
 template <typename SrcOpTy, typename DstOpTy>
 class BinaryOpConversion : public OpConversionPattern<SrcOpTy> {
 public:
   using OpConversionPattern<SrcOpTy>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {

 
 } // namespace
 
 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   populateWithGenerated(ctx, &patterns);
   // clang-format off
   patterns.insert<
+      AnyOpConversion,
       BinaryOpConversion<AddOp, AddIOp>,
       BinaryOpConversion<MulOp, MulIOp>,
       ConstSizeOpConverter,
       GetExtentOpConverter,
       RankOpConverter,
       ShapeOfOpConversion>(ctx);
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }

mlir/lib/Conversion/ShapeToStandard/ShapeToStandardPatterns.td

 // Convert `index_to_size` and `size_to_index` to no-ops as sizes will be
 // represented as indices.
 def IndexToSizeOpConversion : Pat<
     (Shape_IndexToSizeOp $arg),
     (replaceWithValue $arg)>;
 def SizeToIndexOpConversion : Pat<
     (Shape_SizeToIndexOp $arg),
     (replaceWithValue $arg)>;
-

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

func @get_extent_from_extent_tensor(%extents : tensor<?xindex>,
                                     %idx : !shape.size) -> !shape.size {
   // CHECK: %[[RESULT:.*]] = extract_element %[[EXTENTS]][%[[IDX]]] : tensor<?xindex>
   // CHECK: return %[[RESULT]] : index
   %shape = shape.from_extent_tensor %extents : tensor<?xindex>
   %result = shape.get_extent %shape, %idx
   return %result : !shape.size
 }
 
+// -----
+
+// Lower `any` to its first operand.
+// CHECK-LABEL: @any_of_three
+// CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<?xindex>) -> tensor<?xindex>
+func @any_of_three(%a : !shape.shape, %b : !shape.shape, %c : !shape.shape)
+    -> !shape.shape {
+  // CHECK: return %[[A]] : tensor<?xindex>
+  %result = shape.any %a, %b, %c
+  return %result : !shape.shape
+}
+
+// -----
+
+// Lower `any` to its first operand.
+// CHECK-LABEL: @any_of_one
+// CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>) -> tensor<?xindex>
+func @any_of_one(%a : !shape.shape) -> !shape.shape {
+  // CHECK: return %[[A]] : tensor<?xindex>
+  %result = shape.any %a
+  return %result : !shape.shape
+}
+