[MLIR] Add TensorFromElementsOp to Standard ops.

[pifon2a]

Summary

Adds an operation that can create a 1D tensor from a list of values.

[herhut]

Thanks for cleaning this up!

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

1532

Maybe use SameOperandsAndResultElementType?

1556

Would let assemblyFormat = "( $elements ) attr-dict : type($result)"; work?

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

1713

Remove extra '{' '}'.

[pifon2a]

Address the comments.

[herhut]

Thanks.

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

1556

Answer is, it does not because the operand types cannot be deduced with the given traits.

[pifon2a]

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

1532

Thanks, that's a new trait and a useful one. didn't know about it.

1556

no, it wouldn't. It would need to know the elements type to parse it. Can I use assemblyFormat for printing only and a custom parser?

[rriddle]

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

1556

You can do that with TypesMatchWith, this is common:
https://github.com/llvm/llvm-project/blob/0073c293a401774ac96b4b3d27f05e13f379f98e/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td#L484
https://github.com/llvm/llvm-project/blob/0073c293a401774ac96b4b3d27f05e13f379f98e/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td#L1479

[pifon2a]

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

1556

I might be doing smth wrong, but if i have the following op:

mlir
def TensorFromElementsOp : Std_Op<"tensor_from_elements",
    [NoSideEffect, SameOperandsAndResultElementType,
     TypesMatchWith<"value type matches element type of memref",
                    "result", "elements",
                    "$_self.cast<RankedTensorType>().getElementType()">
    ]> {

  let arguments = (ins Variadic<AnyType>:$elements);
  let results = (outs AnyTensor:$result);

  let assemblyFormat = "`(` $elements `)` attr-dict `:` type($result)";
}

then the operands will be resolved with a wrong resolveOperands overload. It should use

ParseResult resolveOperands(ArrayRef<OperandType> operands, Type type,
                            SmallVectorImpl<Value> &result)

instead of

ParseResult resolveOperands(ArrayRef<OperandType> operands,
                              ArrayRef<Type> types, llvm::SMLoc loc,
                              SmallVectorImpl<Value> &result)

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

def ExtractElementOp : Std_Op<"extract_element",
   let hasFolder = 1;
 
   let assemblyFormat = [{
     $aggregate `[` $indices `]` attr-dict `:` type($aggregate)
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// TensorFromElementsOp
+//===----------------------------------------------------------------------===//
+
+def TensorFromElementsOp : Std_Op<"tensor_from_elements",
+    [NoSideEffect, SameOperandsAndResultElementType]> {
+  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 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 hasCanonicalizer = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // FPExtOp
 //===----------------------------------------------------------------------===//
 
 def FPExtOp : CastOp<"fpext">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from floating-point to wider floating-point";
   let description = [{
     Cast a floating-point value to a larger floating-point-typed value.
     The destination type must to be strictly wider than the source type.

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

OpFoldResult ExtractElementOp::fold(ArrayRef<Attribute> operands) {
   // If this is an elements attribute, query the value at the given indices.
   auto elementsAttr = aggregate.dyn_cast<ElementsAttr>();
   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))
+    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.result().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();
+}
+
+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.
+struct ExtractElementFromTensorFromElements
+    : public OpRewritePattern<ExtractElementOp> {
+  using OpRewritePattern<ExtractElementOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExtractElementOp extract,
+                                PatternRewriter &rewriter) const final {
+    if (extract.indices().size() != 1)
+      return failure();
+
+    auto tensor_from_elements =
+        dyn_cast<TensorFromElementsOp>(extract.aggregate().getDefiningOp());
+    if (tensor_from_elements == nullptr)
+      return failure();
+
+    APInt index;
+    if (!matchPattern(*extract.indices().begin(), m_ConstantInt(&index)))
+      return failure();
+    rewriter.replaceOp(extract,
+                       tensor_from_elements.getOperand(index.getZExtValue()));
+    return success();
+  }
+};
+
+} // namespace
+
+void TensorFromElementsOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<ExtractElementFromTensorFromElements>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // FPExtOp
 //===----------------------------------------------------------------------===//
 
 bool FPExtOp::areCastCompatible(Type a, Type b) {
   if (auto fa = a.dyn_cast<FloatType>())
     if (auto fb = b.dyn_cast<FloatType>())
       return fa.getWidth() < fb.getWidth();
   if (auto va = a.dyn_cast<VectorType>())

mlir/test/IR/core-ops.mlir

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

mlir/test/IR/invalid-ops.mlir

func @extract_element_tensor_too_many_indices(%t : tensor<2x3xf32>, %i : index) {
   %0 = "std.extract_element"(%t, %i, %i, %i) : (tensor<2x3xf32>, index, index, index) -> f32
   return
 }
 
 // -----
 
 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
+  %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>
+  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>
   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

 {
   %0 = memref_cast %V : memref<16x16xf32> to memref<?x?xf32>
   %1 = subview %0[0, 0][3, 4][1, 1] : memref<?x?xf32> to memref<3x4xf32, offset:?, strides:[?, 1]>
 
   // CHECK:  subview{{.*}}: memref<16x16xf32> to memref<3x4xf32, #[[map0]]>
   // CHECK:  memref_cast{{.*}}: memref<3x4xf32, #[[map0]]> to memref<3x4xf32, #[[map1]]>
   return %1: memref<3x4xf32, offset:?, strides:[?, 1]>
 }
+
+// -----
+
+// 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>
+  %extracted_element = extract_element %tensor[%c0] : tensor<1xindex>
+  // CHECK: [[ARG]] : index
+  return %extracted_element : index
+}