[MLIR][Shape] Lower `shape.const_shape` to `tensor_from_elements`

[pifon2a]

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

140

nit: no need to check the whole line, you can just drop : tensor<3xindex>

[frgossen]

Rebase

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
 
     // Materialize shape as ranked tensor.
     rewriter.replaceOpWithNewOp<TensorFromElementsOp>(op.getOperation(),
                                                       dimValues);
     return success();
   }
 };
 
+class ConstShapeOpConverter : public OpConversionPattern<ConstShapeOp> {
+public:
+  using OpConversionPattern<ConstShapeOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ConstShapeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    SmallVector<Value, 4> extentOperands;
+    for (auto extent : op.shape()) {
+      extentOperands.push_back(
+          rewriter.create<ConstantIndexOp>(loc, extent.getLimitedValue()));
+    }
+    Value tensor = rewriter.create<TensorFromElementsOp>(loc, extentOperands);
+    Type indexTy = rewriter.getIndexType();
+    Type resultTy = RankedTensorType::get({ShapedType::kDynamicSize}, indexTy);
+    rewriter.replaceOpWithNewOp<TensorCastOp>(op, tensor, resultTy);
+    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>,
+      ConstShapeOpConverter,
       ConstSizeOpConverter,
       RankOpConverter,
       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

 // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
 func @rank(%shape : !shape.shape) -> !shape.size {
   // CHECK-DAG: %[[C0:.*]] = constant 0 : index
   // CHECK-DAG: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]]
   // CHECK-DAG: return %[[RESULT]] : index
   %rank = shape.rank %shape
   return %rank : !shape.size
 }
+
+// -----
+
+// Lower `const_shape` to `tensor_from_elements`.
+// CHECK-LABEL: @const_shape
+// CHECK-SAME: () -> tensor<?xindex>
+func @const_shape() -> !shape.shape {
+  // CHECK: %[[C1:.*]] = constant 1 : index
+  // CHECK: %[[C2:.*]] = constant 2 : index
+  // CHECK: %[[C3:.*]] = constant 3 : index
+  // CHECK: %[[TENSOR3:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]]) : tensor<3xindex>
+  // CHECK: %[[RESULT:.*]] = tensor_cast %[[TENSOR3]] : tensor<3xindex> to tensor<?xindex>
+  // CHECK: return %[[RESULT]] : tensor<?xindex>
+  %shape = shape.const_shape [1, 2, 3]
+  return %shape : !shape.shape
+}