[mlir][shape] Add `shape.from_extents`.

[silvas]

Summary

This is a basic op needed for creating shapes from SSA values
representing the extents.

[jpienaar]

Thanks

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

139

No period at the end please (summary is sentence fragment)

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

212

Seems like we could just have constantShapeReturn function which all the inferReturnType could call [I did a version where it was "trait" that reduced a lot of repetition which I should perhaps clean up a bit - doesn't add anything truly new, just packages]

[silvas]

Address comments.

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

212

I want to avoid that because I specifically want to make this ugly to show the hole in the infra (and eventually make myself fix it :) ). If all result types are buildable types then this should be autogenerated somehow. (as a stretch goal unifying with the type inference logic used by the declarative asm parser)

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

def Shape_ConstSizeOp : Shape_Op<"const_size",
   }];
 
   let arguments = (ins IndexAttr:$value);
   let results = (outs Shape_SizeType:$result);
 
   let assemblyFormat = "attr-dict $value";
 }
 
+def Shape_FromExtentsOp : Shape_Op<"from_extents", [
+    NoSideEffect,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>
+    ]> {
+  let summary = "Creates a shape from extents";
+  let description = [{
+    Creates a shape from multiple SSA values representing the extents of
+    the shape.
+
+    ```mlir
+    // Rank 2 shape.
+    %s0 = shape.from_extents %a, %b
+    // Rank 0 shape.
+    %s1 = shape.from_extents
+    ```
+  }];
+  let arguments = (ins Variadic<Index>:$extents);
+  let results = (outs Shape_ShapeType:$shape);
+
+  let assemblyFormat = "attr-dict $extents";
+
+  let hasFolder = 1;
+}
+
 def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
   let summary = "Creates a shape from a tensor of extents";
   let description = [{
     Creates a shape from a 1D integral tensor of extents. The rank of the
     resulting shape equals the number of elements in the tensor, and the
     extents match the values of the elements.
   }];
 

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

ConstSizeOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
                               SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(SizeType::get(context));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// FromExtentsOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult FromExtentsOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(ShapeType::get(context));
+  return success();
+}
+
+OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
+  if (llvm::any_of(operands, [](Attribute a) { return !a; }))
+    return nullptr;
+  SmallVector<int64_t, 6> extents;
+  for (auto attr : operands)
+    extents.push_back(attr.cast<IntegerAttr>().getInt());
+  Builder builder(getContext());
+  return builder.getI64TensorAttr(extents);
+}
+
+//===----------------------------------------------------------------------===//
 // ShapeOfOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ShapeOfOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                             ValueRange operands, DictionaryAttr attributes,
                             RegionRange regions,
                             SmallVectorImpl<Type> &inferredReturnTypes) {

mlir/test/Dialect/Shape/canonicalize.mlir

 // 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>
 }
+
+// -----
+// Basic case.
+// CHECK-LABEL: func @f()
+func @f() -> !shape.shape {
+  // CHECK: shape.const_shape [3, 5, 11]
+  %e0 = constant 3 : index
+  %e1 = constant 5 : index
+  %e2 = constant 11 : index
+  %ret = shape.from_extents %e0, %e1, %e2
+  return %ret : !shape.shape
+}
+
+// CHECK-LABEL: func @no_fold
+func @no_fold(%arg0: index) -> !shape.shape {
+  // CHECK-NOT: shape.const_shape
+  %e0 = constant 3 : index
+  %ret = shape.from_extents %e0, %arg0
+  return %ret : !shape.shape
+}