[MLIR][Standard] Update `tensor_from_elements` assembly format

[frgossen]

Summary

Remove the redundant parenthesis that are used for none of the other operation
formats.

[herhut]

Thanks for cleaning this up!

[rriddle]

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

1683

Do you know why this is this not using the declarative assembly format?

[frgossen]

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

1683

It seems the assembly format cannot infer the operand types from the the result type.

[frgossen]

Rebase

[rriddle]

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

1683

That is a very common use case, you just need to add the right traits. In this case, TypesMatchWith would do it.

E.g. https://github.com/llvm/llvm-project/blob/df63eedef64d715ce1f31843f7de9c11fe1e597f/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td#L1555

[frgossen]

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

1683

Not sure how that can help here.
The number of elements is variadic and I don't see how to use TypesMatchWith to derive all types simultaneously.

[herhut]

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

1683

See https://github.com/llvm/llvm-project/blob/df63eedef64d715ce1f31843f7de9c11fe1e597f/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td#L730

TypesMatchWith can also iterate over a range of types.

[frgossen]

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

1683

Nice! Now this mini-change evolved into a nice simplification :)

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

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

string description = [{
     tensor_from_elements(i_1, ..., i_N) :  tensor<Nxindex>
     ```
   }];
 
   let arguments = (ins Variadic<AnyType>:$elements);
   let results = (outs AnyTensor:$result);
 
   let skipDefaultBuilders = 1;
-  let builders = [OpBuilder<
-    "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()));
-    }]>];
+  let builders = [
+    OpBuilder<"OpBuilder &b, OperationState &result, ValueRange elements">
+  ];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FPExtOp
 //===----------------------------------------------------------------------===//
 

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.parseLParen() || parser.parseOperandList(elementsOperands) ||
-      parser.parseRParen() || parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColon() || parser.parseType(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 << "tensor_from_elements " << op.elements();
   p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.result().getType();
+  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()));
+}
+
 namespace {
 
 // Canonicalizes the pattern of the form
 //
 // %tensor = "tensor_from_elements(%element) : (i32) -> tensor<1xi32>
 // %extracted_element = extract_element %tensor[%c0] : tensor<1xi32>
 //
 // to just %element.

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

 
 // Lower `shape_of` for statically shaped tensor.
 // CHECK-LABEL: @shape_of_stat
 // CHECK-SAME: (%[[ARG:.*]]: tensor<1x2x3xf32>)
 func @shape_of_stat(%arg : tensor<1x2x3xf32>) {
   // CHECK-DAG: %[[C1:.*]] = constant 1 : index
   // CHECK-DAG: %[[C2:.*]] = constant 2 : index
   // CHECK-DAG: %[[C3:.*]] = constant 3 : index
-  // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]]) : tensor<3xindex>
+  // 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 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>
-  // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements(%[[C1]], %[[C5]], %[[DYN_DIM]]) : tensor<3xindex>
+  // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements %[[C1]], %[[C5]], %[[DYN_DIM]] : tensor<3xindex>
   %shape = shape.shape_of %arg : tensor<1x5x?xf32> -> tensor<?xindex>
   return
 }
 
 // -----
 
 // Convert `rank` to `dim` of the first dimension.
 // CHECK-LABEL: @rank

 
 // Lower `const_shape` to `tensor_from_elements`.
 // CHECK-LABEL: @const_shape
 // CHECK-SAME: () -> tensor<?xindex>
 func @const_shape() -> tensor<?xindex> {
   // CHECK: %[[C1:.*]] = constant 1 : index
   // CHECK: %[[C2:.*]] = constant 2 : index
   // CHECK: %[[C3:.*]] = constant 3 : index
-  // CHECK: %[[TENSOR3:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]])
+  // CHECK: %[[TENSOR3:.*]] = tensor_from_elements %[[C1]], %[[C2]], %[[C3]]
   // 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>
 }
 
 // -----
 

mlir/test/IR/core-ops.mlir

func @extract_element(%arg0: tensor<*xi32>, %arg1 : tensor<4x4xf32>) -> i32 {
   %1 = extract_element %arg1[%c0, %c0] : tensor<4x4xf32>
 
   return %0 : i32
 }
 
 // CHECK-LABEL: func @tensor_from_elements() {
 func @tensor_from_elements() {
   %c0 = "std.constant"() {value = 0: index} : () -> index
-  // CHECK: %0 = tensor_from_elements(%c0) : tensor<1xindex>
-  %0 = tensor_from_elements(%c0) : tensor<1xindex>
+  // CHECK: %0 = tensor_from_elements %c0 : tensor<1xindex>
+  %0 = tensor_from_elements %c0 : tensor<1xindex>
 
   %c1 = "std.constant"() {value = 1: index} : () -> index
-  // CHECK: %1 = tensor_from_elements(%c0, %c1) : tensor<2xindex>
-  %1 = tensor_from_elements(%c0, %c1) : tensor<2xindex>
+  // 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: %2 = tensor_from_elements [[C0_F32]] : tensor<1xf32>
+  %2 = tensor_from_elements %c0_f32 : tensor<1xf32>
 
   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>

mlir/test/IR/invalid-ops.mlir

func @extract_element_tensor_too_few_indices(%t : tensor<2x3xf32>, %i : index) {
   %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}}
   %c0 = constant 0 : i32
-  %0 = tensor_from_elements(%c0) : tensor<*xi32>
+  %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}}
   %c0 = constant 0 : index
-  %0 = tensor_from_elements(%c0) : tensor<2xindex>
+  %0 = tensor_from_elements %c0 : tensor<2xindex>
   return
 }
 
 // -----
 
 func @index_cast_index_to_index(%arg0: index) {
   // expected-error@+1 {{are cast incompatible}}
   %0 = index_cast %arg0: index to index

mlir/test/Transforms/canonicalize.mlir

 }
 
 // -----
 
 // CHECK-LABEL: func @extract_element_from_tensor_from_elements
 func @extract_element_from_tensor_from_elements(%element : index) -> index {
   // CHECK-SAME: ([[ARG:%.*]]: index)
   %c0 = constant 0 : index
-  %tensor = tensor_from_elements(%element) : tensor<1xindex>
+  %tensor = tensor_from_elements %element : tensor<1xindex>
   %extracted_element = extract_element %tensor[%c0] : tensor<1xindex>
   // CHECK: [[ARG]] : index
   return %extracted_element : index
 }