Add type function for ConstShape op.

[jpienaar]

Summary

• Enables inferring return type for ConstShape, takes into account valid return types;
• The compatible return type function could be reused, leaving that for next use refactoring;

[Chia-hungDuan]

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

781–782

I'm thinking if we need to report error here. It seems to me this is to try to infer the return types, any errors may indicate we can't do that, so that is a failed case. What do you think?

789–791

I thinking that do we need to check this or we should have the assumption that it'll be the type we write in ODS? I.e., the size must be 1. Or if we have followed ODS, it's still possible to have size() > 1 case while calling isCompatibleReturnTypes?

In the following check

if (lhs.isa<ShapeType>() || rhs.isa<ShapeType>())
  // Shape type is compatible with all other valid return types.
  return true;

auto lhsTensorType = lhs.cast<TensorType>();
auto rhsTensorType = rhs.cast<TensorType>();

We assume if it's not ShapeType then it will be TensorType, this takes the ODS.

804–816

Can we use verifyCompatibleShapes?

[frgossen]

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

806–810

Can they ever be something other than rank 1?

813

lhsTensorType.isDynamicDim(0) || ...

[jpienaar]

Updating to use verify compatible shape helper

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

let builders = [OpBuilder<(ins "::mlir::Type":$result,
     }]>
   ];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let verifier = [{ return ::verify(*this); }];
 }
 
-def Shape_ConstShapeOp : Shape_Op<"const_shape", [ConstantLike, NoSideEffect]> {
+def Shape_ConstShapeOp : Shape_Op<"const_shape",
+    [ConstantLike, NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Creates a constant shape or extent tensor";
   let description = [{
     Creates a constant shape or extent tensor. The individual extents are given
     as the `shape` attribute. The number of these values equals the shape's
     rank.
 
     ```mlir
     %0 = shape.const_shape [] : !shape.shape
     %1 = shape.const_shape [1, 2, 3] : !shape.shape
-    %2 = shape.const_shape [4, 5, 6] : tensor<?xindex>
+    %2 = shape.const_shape [4, 5, 6] : tensor<3xindex>
     ```
   }];
   let arguments = (ins IndexElementsAttr:$shape);
   let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
   // TODO: Move this to main so that all shape ops implement these.
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let hasFolder = 1;
   let hasCanonicalizer = 1;
+
+  let extraClassDeclaration = [{
+    // InferTypeOpInterface:
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
 }
 
 def Shape_ConstSizeOp : Shape_Op<"const_size", [
     ConstantLike,
     NoSideEffect,
     DeclareOpInterfaceMethods<OpAsmOpInterface>
   ]> {
   let summary = "Creates a constant of type `shape.size`";

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

 
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 

 
 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
 
 void ConstShapeOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                MLIRContext *context) {
   patterns.add<TensorCastConstShape>(context);
 }
 
+LogicalResult mlir::shape::ConstShapeOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  Builder b(context);
+  auto shape = attributes.getAs<DenseIntElementsAttr>("shape");
+  if (!shape)
+    return emitOptionalError(location, "missing shape attribute");
+  inferredReturnTypes.assign({RankedTensorType::get(
+      {static_cast<int64_t>(shape.size())}, b.getIndexType())});
+  return success();
+}
+
+bool mlir::shape::ConstShapeOp::isCompatibleReturnTypes(TypeRange l,
+                                                        TypeRange r) {
+  if (l.size() != 1 || r.size() != 1)
+    return false;
+
+  Type lhs = l.front();
+  Type rhs = r.front();
+
+  if (lhs == rhs)
+    return true;
+
+  if (lhs.isa<ShapeType>() || rhs.isa<ShapeType>())
+    // Shape type is compatible with all other valid return types.
+    return true;
+
+  return succeeded(verifyCompatibleShapes(lhs, rhs));
+}
+
 //===----------------------------------------------------------------------===//
 // CstrBroadcastableOp
 //===----------------------------------------------------------------------===//
 
 void CstrBroadcastableOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
   // Canonicalization patterns have overlap with the considerations during
   // folding in case additional shape information is inferred at some point that
   // does not result in folding.
   patterns.add<CanonicalizeCastExtentTensorOperandsPattern<CstrBroadcastableOp>,
                CstrBroadcastableEqOps,
                RemoveDuplicateOperandsPattern<CstrBroadcastableOp>,
                RemoveEmptyShapeOperandsPattern<CstrBroadcastableOp>>(context);
 }
 
 // Return true if there is exactly one attribute not representing a scalar
 // broadcast.
 static bool hasAtMostSingleNonScalar(ArrayRef<Attribute> attributes) {
   bool nonScalarSeen = false;

mlir/test/Dialect/Shape/invalid.mlir

 
 // -----
 
 func @fn(%arg: !shape.shape) -> !shape.witness {
   // expected-error@+1 {{required at least 2 input shapes}}
   %0 = shape.cstr_broadcastable %arg : !shape.shape
   return %0 : !shape.witness
 }
+
+// -----
+
+// Test that type inference flags the wrong return type.
+
+func @const_shape() {
+  // expected-error@+1 {{'tensor<3xindex>' are incompatible with return type(s) of operation 'tensor<2xindex>'}}
+  %0 = shape.const_shape [4, 5, 6] : tensor<2xindex>
+  return
+}

mlir/test/Dialect/Shape/ops.mlir

func @test_shape_num_elements_unknown() {
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %2 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
 func @const_shape() {
   %0 = shape.const_shape [1, 2, 3] : !shape.shape
   %1 = shape.const_shape [4, 5, 6] : tensor<?xindex>
+  %2 = shape.const_shape [4, 5, 6] : tensor<3xindex>
   return
 }
 
 func @test_shape_num_elements_fixed() {
   %0 = shape.const_shape [1, 57, 92] : !shape.shape
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %3 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return