[MLIR][Shape] Fix lowering of `shape.get_extent`

[frgossen]

Summary

The declarative conversion patterns caused crashes in the asan configuration.
The non-declarative implementation circumvents this.

[pifon2a]

Thank you!

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

104

nit:

if (auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>()) {

[frgossen]

Address nit

[frgossen]

Rebase

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

public:
   matchAndRewrite(ConstSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<ConstantIndexOp>(op.getOperation(),
                                                  op.value().getSExtValue());
     return success();
   }
 };
 
+class GetExtentOpConverter : public OpConversionPattern<GetExtentOp> {
+  using OpConversionPattern<GetExtentOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(GetExtentOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    GetExtentOp::Adaptor transformed(operands);
+
+    // Derive shape extent directly from shape origin if possible.
+    // This circumvents the necessity to materialize the shape in memory.
+    if (auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>()) {
+      rewriter.replaceOpWithNewOp<DimOp>(op, shapeOfOp.arg(),
+                                         transformed.dim());
+      return success();
+    }
+
+    rewriter.replaceOpWithNewOp<ExtractElementOp>(
+        op, rewriter.getIndexType(), transformed.shape(),
+        ValueRange{transformed.dim()});
+    return success();
+  }
+};
+
 class RankOpConverter : public OpConversionPattern<shape::RankOp> {
 public:
   using OpConversionPattern<shape::RankOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     shape::RankOp::Adaptor transformed(operands);

 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   populateWithGenerated(ctx, &patterns);
   // clang-format off
   patterns.insert<
       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

 // represented as indices.
 def IndexToSizeOpConversion : Pat<
     (Shape_IndexToSizeOp $arg),
     (replaceWithValue $arg)>;
 def SizeToIndexOpConversion : Pat<
     (Shape_SizeToIndexOp $arg),
     (replaceWithValue $arg)>;
 
-// Derive shape extent directly from shape origin if possible.
-// This circumvents the necessity to materialize the shape in memory.
-def GetExtentShapeOfConversion : Pat<
-    (Shape_GetExtentOp (Shape_ShapeOfOp $arg), $idx),
-    (Shape_IndexToSizeOp (DimOp $arg, (Shape_SizeToIndexOp $idx))),
-    [],
-    (addBenefit 10)>;
-def GetExtentFromExtentTensorConversion : Pattern<
-    (Shape_GetExtentOp (Shape_FromExtentTensorOp $extents), $idx),
-    [
-      (Shape_SizeToIndexOp:$std_idx $idx),
-      (ExtractElementOp:$std_result $extents, (NativeCodeCall<"ValueRange({$0})"> $std_idx)),
-      (Shape_IndexToSizeOp $std_result)
-    ],
-    [],
-    (addBenefit 10)>;
-