[mlir][Shape] Allow materializing const shapes of any 1D index tensor

[bkramer]

Summary

Limiting it to tensor<?xindex> seems unnecessary. This changes some
test expectations that were relying on other folds.

[herhut]

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

108

Why should this produce a shape.const_shape and not a normal tensor constant? With this setup it is not clear which constant operation will be chosen I think.

[bkramer]

Turned out the thing that motivated this change was a bug elsewhere. dropping.

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

#include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
   // still evolving it makes it simple to start with an unregistered ops and
   // try different variants before actually defining the op.
   allowUnknownOperations();
 }
 
 Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
                                              Attribute value, Type type,
                                              Location loc) {
+  auto rankedType = type.dyn_cast<RankedTensorType>();
   if (type.isa<ShapeType>() ||
-      type == getExtentTensorType(builder.getContext()))
+      (rankedType && rankedType.getRank() == 1 &&
+       rankedType.getElementType() == builder.getIndexType()))
     return builder.create<ConstShapeOp>(loc, type,
                                         value.cast<DenseIntElementsAttr>());
   if (type.isa<SizeType>())
     return builder.create<ConstSizeOp>(loc, type, value.cast<IntegerAttr>());
   if (type.isa<WitnessType>())
     return builder.create<ConstWitnessOp>(loc, type, value.cast<BoolAttr>());
   if (ConstantOp::isBuildableWith(value, type))
     return builder.create<ConstantOp>(loc, type, value);

mlir/test/Dialect/Shape/canonicalize.mlir

 // RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize %s | FileCheck %s
 
 // CHECK-LABEL: func @f
 func @f(%arg0: tensor<2x3x4xf32>) -> tensor<?xindex> {
   // CHECK: shape.const_shape [2, 3, 4] : tensor<?xindex>
   %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
 
 // -----
 
+// CHECK-LABEL: func @f
+func @f(%arg0: tensor<2x3x4xf32>) -> tensor<3xindex> {
+  // CHECK: shape.const_shape [2, 3, 4] : tensor<3xindex>
+  %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<3xindex>
+  return %0 : tensor<3xindex>
+}
+
+// -----
+
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
   // CHECK: shape.const_shape [2, 3] : !shape.shape
   // CHECK: shape.const_shape [4, 5] : !shape.shape
   %c2 = constant 2 : index
   %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
   %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, index) -> (!shape.shape, !shape.shape)
func @f() -> !shape.shape {
   return %0 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> tensor<2xindex> {
-  // CHECK: constant dense<[0, 1]> : tensor<2xindex>
+  // CHECK: shape.const_shape [0, 1] : tensor<2xindex>
   %cs = shape.const_shape [0, 1] : !shape.shape
   %0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
 // -----
 
 // Basic case.
func @fold_assuming_all_single_element(%arg: tensor<?xindex>) {
   return
 }
 
 // -----
 
 // Verify that tensor.cast folding uses the correct type
 // CHECK-LABEL: @fold_tensor.cast_of_const_shape_returned
 func @fold_tensor.cast_of_const_shape_returned(%arg: i1) -> tensor<1xindex> {
-  // CHECK: constant dense<2> : tensor<1xindex>
+  // CHECK: shape.const_shape [2] : tensor<1xindex>
   // CHECK-NOT: tensor.cast
   %0 = shape.const_shape [2] : tensor<?xindex>
   %1 = tensor.cast %0 : tensor<?xindex> to tensor<1xindex>
   return %1 : tensor<1xindex>
 }
 
 // -----