[MLIR][Shape] Lower `shape.shape_of` to standard dialect

[frgossen]

Summary

Lower shape.shape_of to standard dialect.
This lowering supports statically and dynamically shaped tensors.
Support for unranked tensors will be added soon.

[rriddle]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

55

nit: Remove this newline.

[herhut]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

76

The lowering for unranked tensors will need a look, so it will be target scf. How about splitting this pattern in two and make this one fail early for unranked?

[frgossen]

Address comments

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     typename SrcOpTy::Adaptor adaptor(operands);
     rewriter.replaceOpWithNewOp<DstOpTy>(op.getOperation(), adaptor.lhs(),
                                          adaptor.rhs());
     return success();
   }
 };
 
+class ShapeOfOpConversion : public OpConversionPattern<ShapeOfOp> {
+public:
+  using OpConversionPattern<ShapeOfOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    ShapeOfOp::Adaptor transformed(operands);
+    auto loc = op.getLoc();
+    auto tensorVal = transformed.arg();
+    auto tensorTy = tensorVal.getType();
+
+    // For unranked tensors `shape_of` lowers to `scf` and the pattern can be
+    // found in the corresponding pass.
+    if (tensorTy.isa<UnrankedTensorType>())
+      return failure();
+
+    // Build values for individual dimensions.
+    SmallVector<Value, 8> dimValues;
+    auto rankedTensorTy = tensorTy.cast<RankedTensorType>();
+    int64_t rank = rankedTensorTy.getRank();
+    for (int64_t i = 0; i < rank; i++) {
+      if (rankedTensorTy.isDynamicDim(i)) {
+        auto dimVal = rewriter.create<DimOp>(loc, tensorVal, i);
+        dimValues.push_back(dimVal);
+      } else {
+        int64_t dim = rankedTensorTy.getDimSize(i);
+        auto dimVal = rewriter.create<ConstantIndexOp>(loc, dim);
+        dimValues.push_back(dimVal);
+      }
+    }
+
+    // Materialize shape as ranked tensor.
+    rewriter.replaceOpWithNewOp<TensorFromElementsOp>(op.getOperation(),
+                                                      dimValues);
+    return success();
+  }
+};
+
 class ConstSizeOpConverter : public OpConversionPattern<ConstSizeOp> {
 public:
   using OpConversionPattern<ConstSizeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(ConstSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<ConstantIndexOp>(op.getOperation(),

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

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

 // Convert `const_size` to `constant` op.
 // CHECK-LABEL: @size_const
 func @size_const() -> !shape.size {
   %c1 = shape.const_size 1
   return %c1 : !shape.size
 }
 // CHECK: %[[C1:.*]] = constant 1 : index
 // CHECK: return %[[C1]] : index
+
+// -----
+
+// Lower `shape_of` for statically shaped tensor.
+// CHECK-LABEL: @shape_of_stat
+// CHECK-SAME: (%[[ARG:.*]]: tensor<1x2x3xf32>)
+func @shape_of_stat(%arg : tensor<1x2x3xf32>) {
+  // CHECK-DAG: %[[C1:.*]] = constant 1 : index
+  // CHECK-DAG: %[[C2:.*]] = constant 2 : index
+  // CHECK-DAG: %[[C3:.*]] = constant 3 : index
+  // CHECK-DAG: %[[SHAPE:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]]) : tensor<3xindex>
+  %shape = shape.shape_of %arg : tensor<1x2x3xf32>
+  return
+}
+
+// -----
+
+// Lower `shape_of` for dynamically shaped tensor.
+// CHECK-LABEL: @shape_of_dyn
+// CHECK-SAME: (%[[ARG:.*]]: tensor<1x5x?xf32>)
+func @shape_of_dyn(%arg : tensor<1x5x?xf32>) {
+  // CHECK-DAG: %[[C1:.*]] = constant 1 : index
+  // CHECK-DAG: %[[C5:.*]] = constant 5 : index
+  // CHECK-DAG: %[[C2:.*]] = constant 2 : index
+  // CHECK-DAG: %[[DYN_DIM:.*]] = dim %[[ARG]], %[[C2]] : tensor<1x5x?xf32>
+  // CHECK-DAG: %[[SHAPE:.*]] = tensor_from_elements(%[[C1]], %[[C5]], %[[DYN_DIM]]) : tensor<3xindex>
+  %shape = shape.shape_of %arg : tensor<1x5x?xf32>
+  return
+}