[mlir] Fix some edge cases around 0-element TensorFromElementsOp

[silvas]

Summary

This introduces a builder for the more general case that supports zero
elements (where the element type can't be inferred from the ValueRange,
since it might be empty).

Also, fix up some cases in ShapeToStandard lowering that hit this. It
happens very easily when dealing with shapes of 0-D tensors.

The SameOperandsAndResultElementType is redundant with the new
TypesMatchWith and prevented having zero elements.

[jpienaar]

The SameOperandsAndResultElementType is redundant with the new

TypesMatchWith and prevented having zero elements.

So is it only conceptually redudant but the verification differs or is it really redundant? And why is zero elements different?

[frgossen]

Thanks!

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

185

Would it make sense to reuse indexTy (line below) here?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1758

nit: Remove //

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

251

nit: duplicate white line

[ftynse]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1642–1643

Nit: many operations have the types _before_ the values in their constructors, I'd keep it consistent

[silvas]

In D87492#2267795, @jpienaar wrote:

The SameOperandsAndResultElementType is redundant with the new

TypesMatchWith and prevented having zero elements.

So is it only conceptually redudant but the verification differs or is it really redundant? And why is zero elements different?

It is fully redundant, and actually just checks the wrong thing. SameOperandsAndResultElementType assumes >= 1 elements because it gets operands[0] as part of its verification; that's the immediate problem. But SameOperandsAndResultElementType uses getElementTypeOrSelf, so it actually would allow slipping through something like tensor_from_elements %tensor : (tensor<?xf32>) -> tensor<?xf32> (the operands should be f32, not tensor<...f32>) because getElementTypeOrSelf doesn't distinguish tensor<?xf32> and f32.

Our new TypesMatchWith (combined with the ordinary ins/outs checking) is the exact predicate that needs to be checked.

[silvas]

Address comments.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1642–1643

Done. I had to skip the default builders, because it then conflicts with the default builder that takes the result type.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 }
 
 //===----------------------------------------------------------------------===//
 // TensorFromElementsOp
 //===----------------------------------------------------------------------===//
 
 def TensorFromElementsOp : Std_Op<"tensor_from_elements", [
     NoSideEffect,
-    SameOperandsAndResultElementType,
     TypesMatchWith<"operand types match result element type",
                    "result", "elements", "SmallVector<Type, 2>("
                    "$_self.cast<ShapedType>().getDimSize(0), "
                    "$_self.cast<ShapedType>().getElementType())">
   ]> {
   string summary = "tensor from elements operation.";
   string description = [{
     Create a 1D tensor from a range of same-type arguments.
 
     Example:
 
     ```mlir
     tensor_from_elements(i_1, ..., i_N) :  tensor<Nxindex>
     ```
   }];
 
   let arguments = (ins Variadic<AnyType>:$elements);
   let results = (outs 1DTensorOf<[AnyType]>:$result);
 
   let assemblyFormat = "$elements attr-dict `:` type($result)";
 
   // This op is fully verified by its traits.
   let verifier = ?;
 
+  let skipDefaultBuilders = 1;
   let builders = [
+    OpBuilder<"OpBuilder &b, OperationState &result, Type elementType,"
+    "ValueRange elements">,
+    // Special case builder for when `elements` has size >=1.
     OpBuilder<"OpBuilder &b, OperationState &result, ValueRange elements">
   ];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FPExtOp

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

if (op.getType().isa<ShapeType>())
     return failure();
 
   auto loc = op.getLoc();
   SmallVector<Value, 4> extentOperands;
   for (auto extent : op.shape()) {
     extentOperands.push_back(
         rewriter.create<ConstantIndexOp>(loc, extent.getLimitedValue()));
   }
-  Value tensor = rewriter.create<TensorFromElementsOp>(loc, extentOperands);
   Type indexTy = rewriter.getIndexType();
+  Value tensor =
+      rewriter.create<TensorFromElementsOp>(loc, indexTy, extentOperands);
   Type resultTy = RankedTensorType::get({ShapedType::kDynamicSize}, indexTy);
   rewriter.replaceOpWithNewOp<TensorCastOp>(op, tensor, resultTy);
   return success();
 }
 
 namespace {
 class ConstSizeOpConversion : public OpConversionPattern<ConstSizeOp> {
 public:
for (int64_t i = 0; i < rank; i++) {
       } else {
         Value extent =
             rewriter.create<ConstantIndexOp>(loc, rankedTensorTy.getDimSize(i));
         extentValues.push_back(extent);
       }
     }
 
     // Materialize extent tensor.
-    Value staticExtentTensor =
-        rewriter.create<TensorFromElementsOp>(loc, extentValues);
+    Value staticExtentTensor = rewriter.create<TensorFromElementsOp>(
+        loc, rewriter.getIndexType(), extentValues);
     rewriter.replaceOpWithNewOp<TensorCastOp>(op, staticExtentTensor,
                                               op.getType());
     return success();
   }
 
   // Lower to `dynamic_tensor_from_elements` otherwise.
   auto *ctx = rewriter.getContext();
   Value rank = rewriter.create<mlir::RankOp>(loc, tensor);

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

OpFoldResult ExtractElementOp::fold(ArrayRef<Attribute> operands) {
   if (elementsAttr && elementsAttr.isValidIndex(indices))
     return elementsAttr.getValue(indices);
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // TensorFromElementsOp
 //===----------------------------------------------------------------------===//
 
 void TensorFromElementsOp::build(OpBuilder &builder, OperationState &result,
+                                 Type elementType, ValueRange elements) {
+  Type resultTy = RankedTensorType::get({static_cast<int64_t>(elements.size())},
+                                        elementType);
+  result.addOperands(elements);
+  result.addTypes(resultTy);
+}
+
+void TensorFromElementsOp::build(OpBuilder &builder, OperationState &result,
                                  ValueRange elements) {
   assert(!elements.empty() && "expected at least one element");
-  Type resultTy = RankedTensorType::get({static_cast<int64_t>(elements.size())},
-                                        elements.front().getType());
-  build(builder, result, resultTy, elements);
+  build(builder, result, elements.front().getType(), elements);
 }
 
 namespace {
 
 // Canonicalizes the pattern of the form
 //
 // %tensor = "tensor_from_elements(%element) : (i32) -> tensor<1xi32>
 // %extracted_element = extract_element %tensor[%c0] : tensor<1xi32>

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

func @const_shape() -> tensor<?xindex> {
   // CHECK: %[[RESULT:.*]] = tensor_cast %[[TENSOR3]] : tensor<3xindex> to tensor<?xindex>
   // CHECK: return %[[RESULT]] : tensor<?xindex>
   %shape = shape.const_shape [1, 2, 3] : tensor<?xindex>
   return %shape : tensor<?xindex>
 }
 
 // -----
 
+// Lower `const_shape` in the case of rank 0.
+// CHECK-LABEL: func @const_shape_zero_elements
+// CHECK-SAME: () -> tensor<?xindex>
+func @const_shape_zero_elements() -> tensor<?xindex> {
+  // CHECK: %[[TENSOR:.*]] = tensor_from_elements : tensor<0xindex>
+  // CHECK: %[[RESULT:.*]] = tensor_cast %[[TENSOR]] : tensor<0xindex> to tensor<?xindex>
+  // CHECK: return %[[RESULT]] : tensor<?xindex>
+  %shape = shape.const_shape [] : tensor<?xindex>
+  return %shape : tensor<?xindex>
+}
+
+// -----
+
 // Lower `any` to its first operand.
 // CHECK-LABEL: @any_of_three
 // CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<?xindex>) -> tensor<?xindex>
 func @any_of_three(%a : tensor<?xindex>,
                    %b : tensor<?xindex>,
                    %c : tensor<?xindex>) -> tensor<?xindex> {
   // CHECK: return %[[A]] : tensor<?xindex>
   %result = "shape.any"(%a, %b, %c) : (tensor<?xindex>, tensor<?xindex>, tensor<?xindex>) -> tensor<?xindex>
func @shape_of_stat(%arg : tensor<1x2x3xf32>) {
   // CHECK-DAG: %[[C3:.*]] = constant 3 : index
   // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements %[[C1]], %[[C2]], %[[C3]] : tensor<3xindex>
   %shape = shape.shape_of %arg : tensor<1x2x3xf32> -> tensor<?xindex>
   return
 }
 
 // -----
 
+// Lower `shape_of` for 0-D tensor.
+// CHECK-LABEL: @shape_of_zero_d
+// CHECK-SAME: (%[[ARG:.*]]: tensor<f32>)
+func @shape_of_zero_d(%arg : tensor<f32>) {
+  // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements : tensor<0xindex>
+  %shape = shape.shape_of %arg : tensor<f32> -> tensor<?xindex>
+  return
+}
+
+// -----
+
 // Lower `shape_of` for dynamically shaped tensor.
 // CHECK-LABEL: @shape_of_dyn
 // CHECK-SAME: (%[[ARG:.*]]: tensor<1x5x?xf32>)
 func @shape_of_dyn(%arg : tensor<1x5x?xf32>) {
   // CHECK-DAG: %[[C1:.*]] = constant 1 : index
   // CHECK-DAG: %[[C5:.*]] = constant 5 : index
   // CHECK-DAG: %[[C2:.*]] = constant 2 : index
   // CHECK-DAG: %[[DYN_DIM:.*]] = dim %[[ARG]], %[[C2]] : tensor<1x5x?xf32>

mlir/test/IR/core-ops.mlir

func @tensor_from_elements() {
   // CHECK: %1 = tensor_from_elements %c0, %c1 : tensor<2xindex>
   %1 = tensor_from_elements %c0, %c1 : tensor<2xindex>
 
   %c0_f32 = "std.constant"() {value = 0.0: f32} : () -> f32
   // CHECK: [[C0_F32:%.*]] = constant
   // CHECK: %2 = tensor_from_elements [[C0_F32]] : tensor<1xf32>
   %2 = tensor_from_elements %c0_f32 : tensor<1xf32>
 
+  // CHECK: tensor_from_elements : tensor<0xindex>
+  %3 = tensor_from_elements : tensor<0xindex>
+
   return
 }
 
 // CHECK-LABEL: func @tensor_cast(%arg0
 func @tensor_cast(%arg0: tensor<*xf32>, %arg1 : tensor<4x4xf32>, %arg2: tensor<?x?xf32>) {
   // CHECK: %0 = tensor_cast %arg0 : tensor<*xf32> to tensor<?x?xf32>
   %0 = tensor_cast %arg0 : tensor<*xf32> to tensor<?x?xf32>