[MLIR][Standard] Simplify `tensor_from_elements`

[frgossen]

Summary

Define assembly format and add required traits.

[herhut]

Neat!

mlir/test/IR/invalid-ops.mlir

598

Mega-nit: Maybe include the ' at the end, as well.

[frgossen]

Address nit

mlir/test/IR/invalid-ops.mlir

598

Ups

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

let assemblyFormat = [{
     $aggregate `[` $indices `]` attr-dict `:` type($aggregate)
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TensorFromElementsOp
 //===----------------------------------------------------------------------===//
 
-def TensorFromElementsOp : Std_Op<"tensor_from_elements",
-    [NoSideEffect, SameOperandsAndResultElementType]> {
+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 AnyTensor:$result);
+  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, ValueRange elements">
   ];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

if (elementsAttr && elementsAttr.isValidIndex(indices))
     return elementsAttr.getValue(indices);
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // TensorFromElementsOp
 //===----------------------------------------------------------------------===//
 
-static ParseResult parseTensorFromElementsOp(OpAsmParser &parser,
-                                             OperationState &result) {
-  SmallVector<OpAsmParser::OperandType, 4> elementsOperands;
-  Type resultType;
-  if (parser.parseOperandList(elementsOperands) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(resultType))
-    return failure();
-
-  if (parser.resolveOperands(elementsOperands,
-                             resultType.cast<ShapedType>().getElementType(),
-                             result.operands))
-    return failure();
-
-  result.addTypes(resultType);
-  return success();
-}
-
-static void print(OpAsmPrinter &p, TensorFromElementsOp op) {
-  p << "tensor_from_elements " << op.elements();
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.getType();
-}
-
-static LogicalResult verify(TensorFromElementsOp op) {
-  auto resultTensorType = op.result().getType().dyn_cast<RankedTensorType>();
-  if (!resultTensorType)
-    return op.emitOpError("expected result type to be a ranked tensor");
-
-  int64_t elementsCount = static_cast<int64_t>(op.elements().size());
-  if (resultTensorType.getRank() != 1 ||
-      resultTensorType.getShape().front() != elementsCount)
-    return op.emitOpError()
-           << "expected result type to be a 1D tensor with " << elementsCount
-           << (elementsCount == 1 ? " element" : " elements");
-  return success();
-}
-
 void TensorFromElementsOp::build(OpBuilder &builder, OperationState &result,
                                  ValueRange elements) {
   assert(!elements.empty() && "expected at least one element");
-  result.addOperands(elements);
-  result.addTypes(RankedTensorType::get({static_cast<int64_t>(elements.size())},
-                                        *elements.getTypes().begin()));
+  Type resultTy = RankedTensorType::get({static_cast<int64_t>(elements.size())},
+                                        elements.front().getType());
+  build(builder, result, resultTy, 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/IR/invalid-ops.mlir

 func @extract_element_tensor_too_few_indices(%t : tensor<2x3xf32>, %i : index) {
   // expected-error@+1 {{incorrect number of indices for extract_element}}
   %0 = "std.extract_element"(%t, %i) : (tensor<2x3xf32>, index) -> f32 return
 }
 
 // -----
 
 func @tensor_from_elements_wrong_result_type() {
-  // expected-error@+2 {{expected result type to be a ranked tensor}}
+  // expected-error@+2 {{'result' must be 1D tensor of any type values, but got 'tensor<*xi32>'}}
   %c0 = constant 0 : i32
   %0 = tensor_from_elements %c0 : tensor<*xi32>
   return
 }
 
 // -----
 
 func @tensor_from_elements_wrong_elements_count() {
-  // expected-error@+2 {{expected result type to be a 1D tensor with 1 element}}
+  // expected-error@+2 {{1 operands present, but expected 2}}
   %c0 = constant 0 : index
   %0 = tensor_from_elements %c0 : tensor<2xindex>
   return
 }
 
 // -----
 
 func @index_cast_index_to_index(%arg0: index) {