[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

68

nit: Remove this newline.

[herhut]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

89

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

public:
   matchAndRewrite(IndexToSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     IndexToSizeOp::Adaptor transformed(operands);
     rewriter.replaceOp(op.getOperation(), transformed.arg());
     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();
+
+    // Lower `shape_of` for ranked tensors.
+    if (auto rankedTensorTy = tensorTy.dyn_cast<RankedTensorType>()) {
+
+      // Build values for individual dimensions.
+      SmallVector<Value, 8> dimValues;
+      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();
+    }
+
+    // TODO: Implement lowering for unranked tensors.
+    return failure();
+  }
+};
+
 class SizeToIndexOpConversion : public OpConversionPattern<SizeToIndexOp> {
 public:
   using OpConversionPattern<SizeToIndexOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(SizeToIndexOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     SizeToIndexOp::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,
       IndexToSizeOpConversion,
+      ShapeOfOpConversion,
       SizeToIndexOpConversion>(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
+}