}
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;
patterns.insert<
AnyOpConversion,
BinaryOpConversion<AddOp, AddIOp>,
+ ConstShapeOpConverter,
BinaryOpConversion<MulOp, MulIOp>,
GetExtentOpConverter,
RankOpConverter,
// -----
+// 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>