[shape] More constant folding

[silvas]

Summary

• shape split_at
• shape.broadcast
• shape.concat
• shape.to_extent_tensor

Depends On D78754

[jpienaar]

Looks good modulo the same error case comment from last rev, thanks

[silvas]

We currently don't try to constant fold error cases, so that is why we don't try to handle it here.

[silvas]

Update.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

let description = [{
 
     Op has an optional string attribute for the error case where there is no
     broadcastable output shape possible for the given inputs.
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs,
                    OptionalAttr<StrAttr>:$error);
   let results = (outs Shape_ShapeType:$result);
+
+  let hasFolder = 1;
 }
 
 def Shape_ConstShapeOp : Shape_Op<"const_shape",
     [ConstantLike,
      NoSideEffect,
      DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Creates a constant of !shape.shape type.";
   let description = [{
let description = [{
     in the tensor equals the rank of the shape, and the elements equal the
     extents of the shape.
 
     If the shape represents an error, then this op currently aborts the program.
   }];
 
   let arguments = (ins Shape_ShapeType:$input);
   let results = (outs IndexTensor:$result);
+
+  let hasFolder = 1;
 }
 
 def Shape_JoinOp : Shape_Op<"join", []> {
   let summary = "Returns the least general shape.size of its operands";
   let description = [{
     An operation that computes the least general shape of input operands. This
     effectively asserts that corresponding static dimensions are equal. The
     behavior is to match each element of the `shape.type` and propagate the most
let description = [{
     - split_at([4,5,6], index=-4) -> error
 
     Requires:
     - `index` is in the range [-rank(operand),rank(operand)]
   }];
 
   let arguments = (ins Shape_ShapeType:$operand, I32:$index);
   let results = (outs Shape_ShapeType:$head, Shape_ShapeType:$tail);
+  let hasFolder = 1;
 }
 
 def Shape_ConcatOp : Shape_Op<"concat",
     [DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Concatenates two shapes.";
   let description = [{
     Creates a shape whose dimensions consist of first the dimensions from `lhs`
     followed by the dimensions of `rhs`.
 
     Example:
     concat([2,3], [4,5]) -> [2,3,4,5]
     concat([], []) -> []
     concat([], [4,5,6]) -> [4,5,6]
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
   let results = (outs Shape_ShapeType:$result);
+
+  let hasFolder = 1;
 }
 
 #endif // SHAPE_OPS

mlir/lib/Dialect/Shape/CMakeLists.txt

 add_mlir_dialect_library(MLIRShape
   IR/Shape.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
 
   DEPENDS
   MLIRShapeOpsIncGen
   )
 target_link_libraries(MLIRShape
   PUBLIC
+  MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffects
   LLVMSupport
   )

mlir/lib/Dialect/Shape/IR/Shape.cpp

 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Shape/IR/Shape.h"
 
+#include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
case ShapeTypes::ValueShape:
     os << "value_shape";
     return;
   default:
     llvm_unreachable("unexpected 'shape' type kind");
   }
 }
 
 //===----------------------------------------------------------------------===//
+// BroadcastOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
+  if (!operands[0] || !operands[1])
+    return nullptr;
+  auto lhsShape = llvm::to_vector<6>(
+      operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  auto rhsShape = llvm::to_vector<6>(
+      operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  SmallVector<int64_t, 6> resultShape;
+  // If the shapes are not compatible, we can't fold it.
+  // TODO: Fold to an "error".
+  if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
+    return nullptr;
+  Builder builder(getContext());
+  return builder.getI64TensorAttr(resultShape);
+}
+
+//===----------------------------------------------------------------------===//
 // ConstShapeOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, ConstShapeOp &op) {
   p << "shape.const_shape ";
   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"shape"});
   p << "[";
   interleaveComma(op.shape().getValues<int64_t>(), p,
LogicalResult SplitAtOp::inferReturnTypes(
     ArrayRef<NamedAttribute> attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   auto shapeType = ShapeType::get(context);
   inferredReturnTypes.push_back(shapeType);
   inferredReturnTypes.push_back(shapeType);
   return success();
 }
 
+LogicalResult SplitAtOp::fold(ArrayRef<Attribute> operands,
+                              SmallVectorImpl<OpFoldResult> &results) {
+  if (!operands[0] || !operands[1])
+    return failure();
+  auto shapeVec = llvm::to_vector<6>(
+      operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  auto shape = llvm::makeArrayRef(shapeVec);
+  auto splitPoint = operands[1].cast<IntegerAttr>().getInt();
+  // Verify that the split point is in the correct range.
+  // TODO: Constant fold to an "error".
+  int64_t rank = shape.size();
+  if (!(-rank <= splitPoint && splitPoint <= rank))
+    return failure();
+  if (splitPoint < 0)
+    splitPoint += shape.size();
+  Builder builder(operands[0].getContext());
+  results.push_back(builder.getI64TensorAttr(shape.take_front(splitPoint)));
+  results.push_back(builder.getI64TensorAttr(shape.drop_front(splitPoint)));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ConcatOp::inferReturnTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     ArrayRef<NamedAttribute> attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   auto shapeType = ShapeType::get(context);
   inferredReturnTypes.push_back(shapeType);
   return success();
 }
 
+OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
+  if (!operands[0] || !operands[1])
+    return nullptr;
+  auto lhsShape = llvm::to_vector<6>(
+      operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  auto rhsShape = llvm::to_vector<6>(
+      operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  SmallVector<int64_t, 6> resultShape;
+  resultShape.append(lhsShape.begin(), lhsShape.end());
+  resultShape.append(rhsShape.begin(), rhsShape.end());
+  Builder builder(getContext());
+  return builder.getI64TensorAttr(resultShape);
+}
+
+//===----------------------------------------------------------------------===//
+// ToExtentTensorOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ToExtentTensorOp::fold(ArrayRef<Attribute> operands) {
+  if (!operands[0])
+    return nullptr;
+  Builder builder(getContext());
+  auto shape = llvm::to_vector<6>(
+      operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  auto type = RankedTensorType::get({static_cast<int64_t>(shape.size())},
+                                    builder.getIndexType());
+  return DenseIntElementsAttr::get(type, shape);
+}
+
 namespace mlir {
 namespace shape {
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
 
 } // namespace shape
 } // namespace mlir

mlir/test/Dialect/Shape/canonicalize.mlir

-// RUN: mlir-opt -canonicalize <%s | FileCheck %s --dump-input=fail
+// RUN: mlir-opt -split-input-file -canonicalize <%s | FileCheck %s --dump-input=fail
 
 // -----
 // CHECK-LABEL: func @f
 func @f(%arg0: tensor<2x3x4xf32>) -> !shape.shape {
   // CHECK: shape.const_shape [2, 3, 4]
   %0 = "shape.shape_of"(%arg0) : (tensor<2x3x4xf32>) -> !shape.shape
   return %0 : !shape.shape
 }
+
+// -----
+// Basic case.
+// CHECK-LABEL: func @f
+func @f() -> (!shape.shape, !shape.shape) {
+  // CHECK: shape.const_shape [2, 3]
+  // CHECK: shape.const_shape [4, 5]
+  %c2 = constant 2 : i32
+  %0 = "shape.const_shape"() {shape = dense<[2, 3, 4, 5]> : tensor<4xi64>} : () -> !shape.shape
+  %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
+  return %head, %tail : !shape.shape, !shape.shape
+
+}
+
+// -----
+// Negative split point.
+// CHECK-LABEL: func @f
+func @f() -> (!shape.shape, !shape.shape) {
+  // CHECK: shape.const_shape [2, 3, 4]
+  // CHECK: shape.const_shape [5]
+  %c-1 = constant -1 : i32
+  %0 = "shape.const_shape"() {shape = dense<[2, 3, 4, 5]> : tensor<4xi64>} : () -> !shape.shape
+  %head, %tail = "shape.split_at"(%0, %c-1) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
+  return %head, %tail : !shape.shape, !shape.shape
+}
+
+// -----
+// Out of range split point. No folding.
+// CHECK-LABEL: func @f
+func @f() -> (!shape.shape, !shape.shape) {
+  // CHECK: shape.split_at
+  %c5 = constant 5 : i32
+  %0 = "shape.const_shape"() {shape = dense<[2, 3, 4, 5]> : tensor<4xi64>} : () -> !shape.shape
+  %head, %tail = "shape.split_at"(%0, %c5) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
+  return %head, %tail : !shape.shape, !shape.shape
+}
+
+// -----
+// Basic case.
+// CHECK-LABEL: func @f
+func @f() -> !shape.shape {
+  // CHECK: shape.const_shape [7, 2]
+  %0 = shape.const_shape [1, 2]
+  %1 = shape.const_shape [7, 1]
+  %2 = "shape.broadcast"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  return %2 : !shape.shape
+}
+
+// -----
+// Incompatible shapes. No folding.
+// CHECK-LABEL: func @f
+func @f() -> !shape.shape {
+  // CHECK: shape.broadcast
+  %0 = shape.const_shape [2]
+  %1 = shape.const_shape [7]
+  %2 = "shape.broadcast"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  return %2 : !shape.shape
+}
+
+// -----
+// Basic case.
+// CHECK-LABEL: func @f
+func @f() -> !shape.shape {
+  // CHECK: shape.const_shape [0, 1, 2, 3]
+  %lhs = shape.const_shape [0, 1]
+  %rhs = shape.const_shape [2, 3]
+  %0 = "shape.concat"(%lhs, %rhs) : (!shape.shape, !shape.shape) -> !shape.shape
+  return %0 : !shape.shape
+}
+
+// -----
+// Basic case.
+// CHECK-LABEL: func @f
+func @f() -> tensor<2xindex> {
+  // CHECK: constant dense<[0, 1]> : tensor<2xindex>
+  %cs = shape.const_shape [0, 1]
+  %0 = "shape.to_extent_tensor"(%cs) : (!shape.shape) -> tensor<2xindex>
+  return %0 : tensor<2xindex>
+}