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

[pifon2a]

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

150

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

[frgossen]

Rebase

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

LogicalResult ShapeOfOpConversion::matchAndRewrite(
   Value staticExtentTensor =
       rewriter.create<TensorFromElementsOp>(loc, dimValues);
   rewriter.replaceOpWithNewOp<TensorCastOp>(op, staticExtentTensor,
                                             op.getType());
   return success();
 }
 
 namespace {
+class ConstShapeOpConverter : public OpConversionPattern<ConstShapeOp> {
+public:
+  using OpConversionPattern<ConstShapeOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ConstShapeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+} // namespace
+
+LogicalResult ConstShapeOpConverter::matchAndRewrite(
+    ConstShapeOp op, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+
+  // For now, this lowering supports only extent tensors, not `shape.shape`
+  // types.
+  if (op.getType().isa<ShapeType>())
+    return failure();
+
+  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();
+}
+
+namespace {
 class GetExtentOpConverter : public OpConversionPattern<GetExtentOp> {
   using OpConversionPattern<GetExtentOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(GetExtentOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace

 }
 
 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
       AnyOpConversion,
       BinaryOpConversion<AddOp, AddIOp>,
+      ConstShapeOpConverter,
       BinaryOpConversion<MulOp, MulIOp>,
       GetExtentOpConverter,
       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

func @get_extent_from_extent_tensor(%extents : tensor<?xindex>, %idx : index)
   // CHECK: %[[RESULT:.*]] = extract_element %[[EXTENTS]][%[[IDX]]] : tensor<?xindex>
   // CHECK: return %[[RESULT]] : index
   %result = shape.get_extent %extents, %idx : tensor<?xindex>, index -> index
   return %result : index
 }
 
 // -----
 
+// Lower `const_shape` to `tensor_from_elements`.
+// CHECK-LABEL: @const_shape
+// CHECK-SAME: () -> tensor<?xindex>
+func @const_shape() -> tensor<?xindex> {
+  // CHECK: %[[C1:.*]] = constant 1 : index
+  // CHECK: %[[C2:.*]] = constant 2 : index
+  // CHECK: %[[C3:.*]] = constant 3 : index
+  // CHECK: %[[TENSOR3:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]])
+  // CHECK: %[[RESULT:.*]] = tensor_cast %[[TENSOR3]] : tensor<3xindex> to tensor<?xindex>
+  // CHECK: return %[[RESULT]] : tensor<?xindex>
+  %shape = shape.const_shape [1, 2, 3] : tensor<?xindex>
+  return %shape : tensor<?xindex>
+}
+
+// -----
+
 // 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 : tensor<?xindex>,
                    %b : tensor<?xindex>,
                    %c : tensor<?xindex>) -> tensor<?xindex> {
   // CHECK: return %[[A]] : tensor<?xindex>
   %result = "shape.any"(%a, %b, %c) : (tensor<?xindex>, tensor<?xindex>, tensor<?xindex>) -> tensor<?xindex>
   return %result : tensor<?xindex>
 }
 
 // -----
 
 // Lower `any` to its first operand.
 // CHECK-LABEL: @any_of_one
 // CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>) -> tensor<?xindex>
 func @any_of_one(%a : tensor<?xindex>) -> tensor<?xindex> {
   // CHECK: return %[[A]] : tensor<?xindex>
   %result = "shape.any"(%a) : (tensor<?xindex>) -> tensor<?xindex>
   return %result : tensor<?xindex>
 }