[MLIR] Add `from_index`, `to_index`, and `num_elements` to the shape dialect

[frgossen]

Summary

num_elements derives the number of elements from a given shape.
The remaining operations help with conversion between the shape and the standard
dialect types.

[silvas]

mlir/lib/IR/Builders.cpp

57 ↗

(On Diff #264224)

can you split this cleanup of the core infra into a different patch?

[silvas]

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

175

These should be marked NoSideEffect and have folders and InferTypeOpInterface. See https://reviews.llvm.org/D79833 for an example of another op.

[pifon2a]

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

175

I am not sure from_index is enough. In the IR i first read it as if it creates 1D shape with the provided number number of elements. Maybe call this from_index_rank?

Is there some special reason why we have shape.size instead of shape.rank?

[herhut]

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

175

Maybe index_to_size and size_to_index would better convey what this does. We could also have a single operation size_cast in the spirit of the index_cast operation in standard that allows both conversions. That would be more uniform and, if we get to a state where we can have a type trait that would allow to use 'index_cast', it would be easier to rewrite.

[frgossen]

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

175

@silvas NoSideEffect makes perfect sense and I will also add the folders where possible.
Can you explain what the InferTypeOpInterface is good for?
I can copy and adapt the code, of course, but I would like to understand its benefit.
Is the return type of, e.g., Shape_FromIndexOp not already known from the declaration outs Index:$result?

@pifon2a As far as I understand it, shape.size is used to describe also but not only the rank of a tensor.
The shape.reduce example, e.g., already uses it to represent the number of elements.
Until now I thought of shape.size as shape's equivalent to size_t.
The name from/to_index was meant to be analogous to from/to_extent_tensor and index was supposed to refer to std's index.
Would from/to_std_index be a better name?

mlir/lib/IR/Builders.cpp

57 ↗

(On Diff #264224)

It's now in https://reviews.llvm.org/D80111
I will remove the change here.

[silvas]

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

175

I would prefer to have two ops rather than an SizeCast that does everything. That way you can just do isa<SizeToIndexOp> and know what you have. Otherwise, you always have to do isa<SizeCastOp> && op.getResult().getType().isa<SizeType>() or other such nonsense. Personally I think that shape.size_from_index/shape.size_to_index is reads best to me

InferTypeOpInterface avoids needing to redundantly write the result type in the builder calls, since for these ops the result type is already known. (MLIR infra should be smarter about this and do it automatically, but currently that's not the case)

[herhut]

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

175

size_from_index size_to_index works for me, too.

[frgossen]

Rename and constant folding

[frgossen]

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

214

Corrected shape.type to shape.shape and fixed line breaks because the containing line was too long.

[jpienaar]

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

156

Could we keep these sorted?

175

Can you explain what the InferTypeOpInterface is good for?

It is used to generate custom build and verify methods for ops (and pattern rewrites).

I can copy and adapt the code, of course, but I would like to understand its benefit.
Is the return type of, e.g., Shape_FromIndexOp not already known from the declaration > outs Index:$result?

In this case yes, but not yet plumbed through (there is pending item in shape inference doc about this but hasn't been implemented yet - the goal is to have all the shape constraint solving be in one place and be reused). In general these have been type constraints (e.g., Tensor<IntegerType> could be a case so one doesn't have exact shape/type). We should make it easier/better as Sean mentioned.

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

312

It would seem for these,

if (Attribute arg = operands[0]) return arg;
return {};

is simpler.

[silvas]

For future reference, this patch should probably be multiple separate patches:

1. moving ConcatOp
2. adding size_to_index and size_from_index
3. adding num_elements
4. Tidying up docs of Shape_JoinOp and Shape_ReduceOp and Shape_ConstSizeOp
5. the change to StandardTypes.h

That helps keep each review focused which results in much better review feedback. It also helps people retroactively working through the code history to understand the development process, or folks trying.

I would appreciate if you do it for this patch already; I think there's more feedback to be gotten that is being masked by the mix of stuff happening in this patch. At the very least the commit description needs to be updated to incorporate the fact that it is doing all this stuff.

mlir/test/Dialect/Shape/canonicalize.mlir

119

For this specific test, you might want to specifically check that the value is returned directly (otherwise the test also passes if no transformation is made)

128

You don't seem to use these captured FileCheck variables.

[frgossen]

As suggested by @silvas , this revision is now split up into multiple ones:

• [MLIR] Cosmetic change: https://reviews.llvm.org/D80275
• [MLIR] Move ConcatOp to its lexicographic position: https://reviews.llvm.org/D80277
• [MLIR] Tidy up documentation for Shape_JoinOp, Shape_ReduceOp, and Shape_ConstSizeOp: https://reviews.llvm.org/D80278
• [MLIR] Add index_to_size and size_to_index to the shape dialect: https://reviews.llvm.org/D80280?vs=on&id=265167
• [MLIR] Add num_elements to the shape dialect: https://reviews.llvm.org/D80281
• [MLIR] Fix operand type in from_extent_tensor in the shape dialect: https://reviews.llvm.org/D80283

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

156

See https://reviews.llvm.org/D80280

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

312

See https://reviews.llvm.org/D80280

mlir/test/Dialect/Shape/canonicalize.mlir

119

See https://reviews.llvm.org/D80280

128

See https://reviews.llvm.org/D80280

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

def Shape_ConstShapeOp : Shape_Op<"const_shape",
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let hasFolder = 1;
 }
 
 def Shape_ConstSizeOp : Shape_Op<"const_size",
     [ConstantLike,
      NoSideEffect,
      DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
-  let summary = "Creates a constant of !shape.size type.";
+  let summary = "Creates a constant of type `!shape.size`";
   let description = [{
-    Creates a !shape.size type representing the constant size given by `value`.
+    Creates a `!shape.size` type representing the constant size given by `value`.
 
     ```mlir
     %x = shape.const_size 10
     ```
   }];
 
   let arguments = (ins IndexAttr:$value);
   let results = (outs Shape_SizeType:$result);
 
   let assemblyFormat = "attr-dict $value";
+
+  let hasFolder = 1;
+}
+
+def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [
+    NoSideEffect,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Converts a standard index to a shape size";
+  let description = [{
+    Converts a standard index to a `shape.size`.
+    This operation and its inverse, `size_to_index`, facilitate index conversion
+    between the standard and the shape dialect.
+    Expects an `index` and returns a `shape.size` value.
+  }];
+
+  let arguments = (ins Index:$arg);
+  let results = (outs Shape_SizeType:$result);
+
+  let assemblyFormat = "attr-dict $arg";
+
+  let hasFolder = 1;
+}
+
+def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [
+    NoSideEffect,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Casts between index types of the shape and standard dialect";
+  let description = [{
+    Converts a `shape.size` to a standard index.
+    This operation and its inverse, `index_to_size`, facilitate index conversion
+    between the standard and the shape dialect.
+    Expects a `shape.size` and returns an `index` value.
+  }];
+
+  let arguments = (ins Shape_SizeType:$arg);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = "attr-dict $arg";
+
+  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.
   }];
 
-  let arguments = (ins I32Tensor:$input);
+  let arguments = (ins IndexTensor:$input);
   let results = (outs Shape_ShapeType:$result);
 }
 
 def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", []> {
   let summary = "Creates a dimension tensor from a shape";
   // TODO: Think more about the error situation. Perhaps factor out the
   // error detection into a separate op so downstream consumers can control
   // their error behavior. Then this op would assume that the input has
def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", []> {
   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
-    restrictive information, returning an invalid shape if there are
+    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.shape` and propagate the
+    most restrictive information, returning an invalid shape if there are
     contradictory requirements. E.g., using pseudo code
 
     ```
     shape.join([*], [*]) -> [*]
     shape.join([*], [1, ?]) -> [1, ?]
     shape.join([1, 2], [1, ?]) -> [1, 2]
     shape.join([*], [1, 2]) -> [1, 2]
     shape.join([], []) -> []
     shape.join([], [*]) -> []
     shape.join([], [?, ?]) -> [invalid]
     shape.join([1, ?], [2, ?, ?]) -> [invalid]
     ```
 
     `shape.join` also allows specifying an optional error string, that may be
     used to return an error to the user upon mismatch of dimensions.
 
     ```mlir
-    %c = shape.join %a, %b, error="<reason>" : !shape.type
+    %c = shape.join %a, %b, error="<reason>" : !shape.shape
     ```
   }];
 
   let arguments = (ins Shape_ShapeOrSizeType:$arg0, Shape_ShapeOrSizeType:$arg1,
                    OptionalAttr<StrAttr>:$error);
   let results = (outs Shape_ShapeOrSizeType:$result);
 }
 
 def Shape_MulOp : Shape_Op<"mul", [SameOperandsAndResultType]> {
   let summary = "Multiplication of sizes";
   let description = [{
     Multiplies two valid sizes as follows:
     - lhs * rhs = unknown if either lhs or rhs unknown;
     - lhs * rhs = (int)lhs * (int)rhs if both known;
   }];
 
   let arguments = (ins Shape_SizeType:$lhs, Shape_SizeType:$rhs);
   let results = (outs Shape_SizeType:$result);
 }
 
+def Shape_NumElementsOp : Shape_Op<"num_elements", [
+    NoSideEffect,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+
+  let summary = "Returns the number of elements for a given shape";
+  let description = [{
+    Returns the number of elements for a given shape which is the product of its
+    dimensions.
+    A tensor of the given shape will hold this many elements.
+    Expects a `shape.shape` and returns a `shape.size` value.
+  }];
+
+  let arguments = (ins Shape_ShapeType:$shape);
+  let results = (outs Shape_SizeType:$result);
+
+  let assemblyFormat = "attr-dict $shape";
+
+  let hasFolder = 1;
+}
+
 def Shape_ReduceOp : Shape_Op<"reduce", []> {
   let summary = "Returns an expression reduced over a shape";
   let description = [{
     An operation that takes as input a shape, number of initial values and has a
     region/function that is applied repeatedly for every dimension of the shape.
 
     Conceptually this op performs the following reduction:
 
     ```
     res[] = init;
     for (int i = 0, e = shape.rank(); i != e; ++i) {
       res = fn(i, shape[i], res[0], ..., res[n]);
     }
     ```
 
     Where fn is provided by the user and the result of the reduce op is the
     last computed output of the reduce function. As an example, computing the
     number of elements
 
     ```mlir
-    func @shape_num_elements(%shape : !shape.type) -> !shape.size {
+    func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
       %0 = "shape.constant_dim"() {value = 1 : i32} : () -> !shape.size
       %1 = "shape.reduce"(%shape, %0) ( {
         ^bb0(%index: i32, %dim: !shape.size, %lci: !shape.size):
           %acc = "shape.mul"(%lci, %dim) :
             (!shape.size, !shape.size) -> !shape.size
           shape.yield %acc : !shape.size
-        }) : (!shape.type, !shape.size) -> (!shape.size)
+        }) : (!shape.shape, !shape.size) -> (!shape.size)
       return %1 : !shape.size
     }
     ```
 
     If the shape is unranked, then the results of the op is also unranked.
   }];
 
   let arguments = (ins Shape_ShapeType:$shape, Variadic<AnyType>:$args);

mlir/include/mlir/IR/StandardTypes.h

public:
   static bool classof(Type type) {
     return type.getKind() == StandardTypes::RankedTensor ||
            type.getKind() == StandardTypes::UnrankedTensor;
   }
 };
 
 //===----------------------------------------------------------------------===//
 // RankedTensorType
+//===----------------------------------------------------------------------===//
 
 /// Ranked tensor types represent multi-dimensional arrays that have a shape
 /// with a fixed number of dimensions. Each shape element can be a non-negative
 /// integer or unknown (represented by -1).
 class RankedTensorType
     : public Type::TypeBase<RankedTensorType, TensorType,
                             detail::RankedTensorTypeStorage> {
 public:

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

OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
   // TODO: Fold to an "error".
   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
     return nullptr;
   Builder builder(getContext());
   return builder.getI64TensorAttr(resultShape);
 }
 
 //===----------------------------------------------------------------------===//
+// ConcatOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+ConcatOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
+                           ValueRange operands, DictionaryAttr 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);
+}
+
+//===----------------------------------------------------------------------===//
 // 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,
static ParseResult parseConstShapeOp(OpAsmParser &parser,
   }
   Builder &builder = parser.getBuilder();
   result.addAttribute("shape", builder.getI64TensorAttr(ints));
 
   result.types.push_back(ShapeType::get(builder.getContext()));
   return success();
 }
 
-OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shape(); }
+OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
 
 LogicalResult
 ConstShapeOp::inferReturnTypes(MLIRContext *context,
                                Optional<Location> location, ValueRange operands,
                                DictionaryAttr attributes, RegionRange regions,
                                SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ConstSizeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ConstSizeOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
                               SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(SizeType::get(context));
   return success();
 }
 
+OpFoldResult ConstSizeOp::fold(ArrayRef<Attribute>) { return valueAttr(); }
+
+//===----------------------------------------------------------------------===//
+// IndexToSizeOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
+
+  // Fold only when argument constant.
+  Attribute arg = operands[0];
+  if (!arg)
+    return {};
+
+  // Constant values of both types, `shape.size` and `index`, are represented as
+  // `IntegerAttr`s which makes constant folding simple.
+  return arg;
+}
+
+LogicalResult IndexToSizeOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(SizeType::get(context));
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// NumElementsOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
+
+  // Fold only when argument constant.
+  Attribute shape = operands[0];
+  if (!shape)
+    return {};
+
+  APInt product(64, 1);
+  for (auto value : shape.cast<DenseIntElementsAttr>())
+    product *= value;
+  Builder builder(getContext());
+  return builder.getIndexAttr(product.getLimitedValue());
+}
+
+LogicalResult NumElementsOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(SizeType::get(context));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ShapeOfOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ShapeOfOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                             ValueRange operands, DictionaryAttr attributes,
                             RegionRange regions,
                             SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
   auto type = getOperand().getType().dyn_cast<ShapedType>();
   if (!type || !type.hasStaticShape())
     return nullptr;
   Builder builder(getContext());
   return builder.getI64TensorAttr(type.getShape());
 }
 
 //===----------------------------------------------------------------------===//
+// SizeToIndexOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
+
+  // Fold only when argument constant.
+  Attribute arg = operands[0];
+  if (!arg)
+    return {};
+
+  // Constant values of both types, `shape.size` and `index`, are represented as
+  // `IntegerAttr`s which makes constant folding simple.
+  return arg;
+}
+
+LogicalResult SizeToIndexOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(IndexType::get(context));
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // SplitAtOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 SplitAtOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                             ValueRange operands, DictionaryAttr attributes,
                             RegionRange regions,
                             SmallVectorImpl<Type> &inferredReturnTypes) {
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, DictionaryAttr 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>(

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>
 }
+
+// -----
+// Cast constant size to index and fold it away.
+// CHECK-LABEL: func @const_size_to_index
+func @const_size_to_index() -> index {
+  // CHECK-NOT: shape.index_cast
+  %cs = shape.const_size 123
+  // CHECK: constant 123 : index
+  %ci = shape.size_to_index %cs
+  return %ci : index
+}
+
+// -----
+// Cast constant index to size and fold it away.
+// CHECK-LABEL: func @const_index_to_size
+func @const_index_to_size() -> !shape.size {
+  // CHECK-NOT: index_cast
+  %ci = constant 123 : index
+  // CHECK: shape.const_size 123
+  %cs = shape.index_to_size %ci
+  return %cs : !shape.size
+}
+
+// -----
+// Cast constant index to size, then back, and fold it away.
+// CHECK-LABEL: func @const_index_to_size_to_index
+func @const_index_to_size_to_index() -> index {
+  // CHECK-NOT: shape.index_cast
+  %ci0 = constant 123 : index
+  %cs0 = shape.index_to_size %ci0
+  // CHECK: constant 123 : index
+  %ci1 = shape.size_to_index %cs0
+  return %ci1 : index
+}
+
+// -----
+// No folding.
+// CHECK-LABEL: func @nonfoldable_size_to_index
+func @nonfoldable_size_to_index(%cs : !shape.size) -> index {
+  // CHECK: %[[CI:.*]] = shape.size_to_index %[[CS:.*]]
+  %ci = shape.size_to_index %cs
+  return %ci : index
+}
+
+// -----
+// No folding.
+// CHECK-LABEL: func @nonfoldable_index_to_size
+func @nonfoldable_index_to_size(%ci : index) -> !shape.size {
+  // CHECK: %[[CS:.*]] = shape.index_to_size %[[CI:.*]]
+  %cs = shape.index_to_size %ci
+  return %cs : !shape.size
+}
+
+// -----
+// Fold number of elements computation.
+// CHECK-LABEL: func @num_elements
+func @num_elements() -> !shape.size {
+  // CHECK-NOT: shape.const_shape
+  %shape = shape.const_shape [4, 5, 6]
+  // CHECK-NOT: shape.num_elements
+  %num_elements = shape.num_elements %shape
+  // CHECK: %[[NUM:.*]] = shape.const_size 120
+  // CHECK-NEXT: return %[[NUM]] : !shape.size
+  return %num_elements : !shape.size
+}
+
+// -----
+// No folding.
+// CHECK-LABEL: func @nonfoldable_num_elements
+func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size {
+  // CHECK-NOT: shape.const_{{.*}}
+  %num_elements = shape.num_elements %shape
+  return %num_elements : !shape.size
+}