[mlir] Extend `tensor.from_elements` to support N-D case.

[pifon2a]

Summary

RFC: https://llvm.discourse.group/t/rfc-extend-tensor-fromelementsop-to-n-d/4715

[pifon2a]

update

[pifon2a]

update

[herhut]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

322

Can you also mention how the order of elements relates to elements in the tensor?

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

198

Debugging leftover?

199

Could this use the overload that just takes the values?

572

Here, too.

573

Could this use the overload that just takes extenTensor?

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

413

This needs to be the stride. So you can compute flatIndex *= tensorType.getDimSize(i) and then do flatIndex += index.getSExtValue.

mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp

85

Would it be easier to always allocate a 1d memref, write the values and then reshape it into the final shape?

[pifon2a]

Address the comments.

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

199

only if we are absolutely sure that the elements are not empty. Unfortunately, there is this ugly use case:

%0 = tensor.from_elements : tensor<0xindex>

573

There is a builder that does it. But we have to be sure that extentValues are not empty.

mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp

85

It would be the same number of stores + additional reshape that might break some canonicalizations. I think the best variant would be to have a memref.from_elements and then we could lower it to llvm creating a necessary descriptor and linearly populating the underlying buffer.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

 //===----------------------------------------------------------------------===//
 // FromElementsOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_FromElementsOp : Tensor_Op<"from_elements", [
     NoSideEffect,
     TypesMatchWith<"operand types match result element type",
                    "result", "elements", "SmallVector<Type, 2>("
-                   "$_self.cast<ShapedType>().getDimSize(0), "
+                   "$_self.cast<ShapedType>().getNumElements(), "
                    "$_self.cast<ShapedType>().getElementType())">
   ]> {
   string summary = "tensor from elements operation.";
   string description = [{
-    Create a 1D tensor from a range of same-type arguments.
+    Create a N-D tensor from a range of same-type arguments. The number of
+    provided `elements` should equal to the number of the elements in the
+    result type.
 
     Example:
 
     ```mlir
-    tensor.from_elements i_1, ..., i_N :  tensor<Nxindex>
+    tensor.from_elements i_1, ..., i_{N * M} :  tensor<NxMxindex>
     ```
   }];
 
   let arguments = (ins Variadic<AnyType>:$elements);
-  let results = (outs 1DTensorOf<[AnyType]>:$result);
+  let results = (outs AnyStaticShapeTensor:$result);
 
   let assemblyFormat = "$elements attr-dict `:` type($result)";
 
   // This op is fully verified by its traits.
   let verifier = ?;
 
   let skipDefaultBuilders = 1;
   let builders = [
-    OpBuilder<(ins "Type":$elementType, "ValueRange":$elements)>,
+    OpBuilder<(ins "Type":$resultType, "ValueRange":$elements)>,
     // Special case builder for when `elements` has size >=1.
     OpBuilder<(ins "ValueRange":$elements)>
   ];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 

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.getShape()) {
     extentOperands.push_back(
         rewriter.create<arith::ConstantIndexOp>(loc, extent.getLimitedValue()));
   }
-  Type indexTy = rewriter.getIndexType();
+  Type resultTy =
+      RankedTensorType::get({op.getShape().size()}, rewriter.getIndexType());
+  llvm::errs() << "da funck2\n";
   Value tensor =
-      rewriter.create<tensor::FromElementsOp>(loc, indexTy, extentOperands);
-  Type resultTy = RankedTensorType::get({op.getShape().size()}, indexTy);
+      rewriter.create<tensor::FromElementsOp>(loc, resultTy, extentOperands);
   rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultTy, tensor);
   return success();
 }
 
 namespace {
 class ConstSizeOpConversion : public OpConversionPattern<ConstSizeOp> {
 public:
   using OpConversionPattern<ConstSizeOp>::OpConversionPattern;
for (int64_t i = 0; i < rank; i++) {
       } else {
         Value extent = rewriter.create<arith::ConstantIndexOp>(
             loc, rankedTensorTy.getDimSize(i));
         extentValues.push_back(extent);
       }
     }
 
     // Materialize extent tensor.
+    llvm::errs() << "da funck1\n";
     Value staticExtentTensor = rewriter.create<tensor::FromElementsOp>(
-        loc, rewriter.getIndexType(), extentValues);
+        loc, RankedTensorType::get({rank}, rewriter.getIndexType()),
+        extentValues);
     rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(),
                                                 staticExtentTensor);
     return success();
   }
 
   // Lower to `tensor.generate` otherwise.
   auto *ctx = rewriter.getContext();
   Value rank = rewriter.create<tensor::RankOp>(loc, tensor);

mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp

 static Value sourceMaterializationCallback(OpBuilder &builder, Type type,
                                            ValueRange inputs, Location loc) {
   assert(inputs.size() == 1);
   if (inputs[0].getType().isa<TensorType>())
     return nullptr;
 
   // A detensored value is converted back by creating a new tensor from its
   // element(s).
-  auto createNewTensorOp = builder.create<tensor::FromElementsOp>(
-      loc, inputs[0].getType(), inputs[0]);
+  auto createNewTensorOp =
+      builder.create<tensor::FromElementsOp>(loc, inputs[0]);
 
   // FromElementsOp results in a tensor<1xdtype>, we need to reshape that to
   // a tensor<dtype> instead.
   return builder.create<tensor::CollapseShapeOp>(
       loc, type, createNewTensorOp, ArrayRef<ReassociationExprs>{});
 }
 
 namespace {

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

OpFoldResult ExtractOp::fold(ArrayRef<Attribute> operands) {
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // FromElementsOp
 //===----------------------------------------------------------------------===//
 
 void FromElementsOp::build(OpBuilder &builder, OperationState &result,
-                           Type elementType, ValueRange elements) {
-  Type resultTy = RankedTensorType::get({static_cast<int64_t>(elements.size())},
-                                        elementType);
+                           Type resultType, ValueRange elements) {
   result.addOperands(elements);
-  result.addTypes(resultTy);
+  result.addTypes(resultType);
 }
 
 void FromElementsOp::build(OpBuilder &builder, OperationState &result,
                            ValueRange elements) {
   assert(!elements.empty() && "expected at least one element");
-  build(builder, result, elements.front().getType(), elements);
+  Type resultType = RankedTensorType::get(
+      {static_cast<int64_t>(elements.size())}, elements.front().getType());
+  build(builder, result, resultType, elements);
 }
 
 OpFoldResult FromElementsOp::fold(ArrayRef<Attribute> operands) {
   if (!llvm::is_contained(operands, nullptr))
     return DenseElementsAttr::get(getType(), operands);
   return {};
 }
 
 namespace {
 
 // Canonicalizes the pattern of the form
 //
 // %tensor = tensor.from_elements(%element) : (i32) -> tensor<1xi32>
 // %extracted_element = tensor.extract %tensor[%c0] : tensor<1xi32>
 //
 // to just %element.
 struct ExtractElementFromTensorFromElements
     : public OpRewritePattern<tensor::ExtractOp> {
   using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::ExtractOp extract,
                                 PatternRewriter &rewriter) const final {
-    if (extract.indices().size() != 1)
-      return failure();
-
     auto tensorFromElements = extract.tensor().getDefiningOp<FromElementsOp>();
-    if (tensorFromElements == nullptr)
+    if (!tensorFromElements)
       return failure();
-
+    auto tensorType = tensorFromElements.getType().cast<RankedTensorType>();
+    auto rank = tensorType.getRank();
+    SmallVector<APInt, 3> indices(rank);
+    int64_t flatIndex = 0;
+    for (int i = 0; i < rank; ++i) {
       APInt index;
-    if (!matchPattern(*extract.indices().begin(), m_ConstantInt(&index)))
+      if (!matchPattern(extract.indices()[i], m_ConstantInt(&index)))
         return failure();
+      flatIndex += i == rank - 1
+                       ? index.getSExtValue()
+                       : index.getSExtValue() * tensorType.getDimSize(i + 1);
+    }
     // Prevent out of bounds accesses. This can happen in invalid code that will
     // never execute.
-    if (tensorFromElements->getNumOperands() <= index.getZExtValue() ||
-        index.getSExtValue() < 0)
+    if (tensorFromElements->getNumOperands() <= flatIndex || flatIndex < 0)
       return failure();
-    rewriter.replaceOp(extract,
-                       tensorFromElements.getOperand(index.getZExtValue()));
+    rewriter.replaceOp(extract, tensorFromElements.getOperand(flatIndex));
     return success();
   }
 };
 
 } // namespace
 
 void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {

mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp

 #include "PassDetail.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
 struct BufferizeCastOp : public OpConversionPattern<tensor::CastOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult

 
 struct BufferizeFromElementsOp
     : public OpConversionPattern<tensor::FromElementsOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tensor::FromElementsOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    int numberOfElements = op.elements().size();
-    auto resultType = MemRefType::get(
-        {numberOfElements}, op.getType().cast<TensorType>().getElementType());
-    Value result = rewriter.create<memref::AllocOp>(op.getLoc(), resultType);
-    for (auto element : llvm::enumerate(op.elements())) {
-      Value index =
-          rewriter.create<arith::ConstantIndexOp>(op.getLoc(), element.index());
-      rewriter.create<memref::StoreOp>(op.getLoc(), element.value(), result,
-                                       index);
+    Location loc = op.getLoc();
+    auto tensorType = op.getType().cast<RankedTensorType>();
+    auto shape = tensorType.getShape();
+
+    // Allocate a buffer for the result.
+    auto resultType =
+        MemRefType::get(tensorType.getShape(), tensorType.getElementType());
+    Value buffer = rewriter.create<memref::AllocOp>(loc, resultType);
+
+    // Case: tensor<0xelem_type>.
+    if (op.elements().empty()) {
+      rewriter.replaceOp(op, {buffer});
+      return success();
     }
-    rewriter.replaceOp(op, {result});
+
+    // Case: tensor<elem_type>.
+    if (shape.empty()) {
+      rewriter.create<memref::StoreOp>(loc, op.elements().front(), buffer);
+      rewriter.replaceOp(op, {buffer});
       return success();
     }
+
+    // Create constants for the range of possible indices [0, max{shape_i}).
+    auto maxDim = *std::max_element(shape.begin(), shape.end());
+    SmallVector<Value, 2> constants;
+    constants.reserve(maxDim);
+    for (int i = 0; i < maxDim; ++i)
+      constants.push_back(rewriter.create<arith::ConstantIndexOp>(loc, i));
+
+    // Traverse all `elements` and create `memref.store` ops.
+    ImplicitLocOpBuilder b(loc, rewriter);
+    auto element_it = adaptor.elements().begin();
+    SmallVector<Value, 2> indices(tensorType.getRank(), constants[0]);
+    CreateStores(/*dim=*/0, buffer, shape, constants, element_it, indices, b);
+
+    rewriter.replaceOp(op, {buffer});
+    return success();
+  }
+
+private:
+  // Implements backtracking to traverse indices of the output buffer while
+  // iterating over op.elements().
+  void CreateStores(int dim, Value buffer, ArrayRef<int64_t> shape,
+                    ArrayRef<Value> constants, ValueRange::iterator &element_it,
+                    SmallVectorImpl<Value> &indices,
+                    ImplicitLocOpBuilder b) const {
+    if (dim == shape.size() - 1) {
+      for (int i = 0; i < shape.back(); ++i) {
+        indices.back() = constants[i];
+        b.create<memref::StoreOp>(*element_it, buffer, indices);
+        ++element_it;
+      }
+      return;
+    }
+    for (int i = 0; i < shape[dim]; ++i) {
+      indices[dim] = constants[i];
+      CreateStores(dim + 1, buffer, shape, constants, element_it, indices, b);
+    }
+  }
 };
 
 struct BufferizeGenerateOp : public OpConversionPattern<tensor::GenerateOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tensor::GenerateOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {

mlir/test/Dialect/Tensor/bufferize.mlir

 // CHECK:           %[[RET:.*]] = memref.load %[[MEMREF]][%[[IDX]]] : memref<?xf32>
 // CHECK:           return %[[RET]] : f32
 // CHECK:         }
 func @tensor.extract(%arg0: tensor<?xf32>, %arg1: index) -> f32 {
   %0 = tensor.extract %arg0[%arg1] : tensor<?xf32>
   return %0 : f32
 }
 
-// CHECK-LABEL:   func @tensor.from_elements(
+// CHECK-LABEL:   func @tensor.from_elements_no_elements() -> tensor<0xindex> {
+// CHECK:           %[[MEMREF:.*]] = memref.alloc() : memref<0xindex>
+// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
+// CHECK:           return %[[RET]] : tensor<0xindex>
+func @tensor.from_elements_no_elements() -> tensor<0xindex> {
+  %0 = tensor.from_elements : tensor<0xindex>
+  return %0 : tensor<0xindex>
+}
+
+// CHECK-LABEL:   func @tensor.from_elements_0d(
+// CHECK-SAME:        %[[ELEM0:.*]]: index) -> tensor<index> {
+// CHECK:           %[[MEMREF:.*]] = memref.alloc() : memref<index>
+// CHECK:           store %[[ELEM0]], %[[MEMREF]]
+// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
+// CHECK:           return %[[RET]] : tensor<index>
+func @tensor.from_elements_0d(%arg0: index) -> tensor<index> {
+  %0 = tensor.from_elements %arg0 : tensor<index>
+  return %0 : tensor<index>
+}
+
+// CHECK-LABEL:   func @tensor.from_elements_1d(
 // CHECK-SAME:                               %[[ELEM0:.*]]: index,
 // CHECK-SAME:                               %[[ELEM1:.*]]: index) -> tensor<2xindex> {
-// CHECK:           %[[MEMREF:.*]] = memref.alloc()
+// CHECK:           %[[MEMREF:.*]] = memref.alloc() : memref<2xindex>
 // CHECK:           %[[C0:.*]] = arith.constant 0 : index
-// CHECK:           store %[[ELEM0]], %[[MEMREF]][%[[C0]]]
 // CHECK:           %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           store %[[ELEM0]], %[[MEMREF]][%[[C0]]]
 // CHECK:           store %[[ELEM1]], %[[MEMREF]][%[[C1]]]
 // CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
 // CHECK:           return %[[RET]] : tensor<2xindex>
-func @tensor.from_elements(%arg0: index, %arg1: index) -> tensor<2xindex> {
+func @tensor.from_elements_1d(%arg0: index, %arg1: index) -> tensor<2xindex> {
   %0 = tensor.from_elements %arg0, %arg1 : tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
+// CHECK-LABEL: func @tensor.from_elements_2d(
+// CHECK-SAME:      %[[ELEM0:.*]]: index, %[[ELEM1:.*]]: index)
+// CHECK-SAME:      -> tensor<3x2xindex> {
+// CHECK:         %[[MEMREF:.*]] = memref.alloc() : memref<3x2xindex>
+// CHECK:         %[[C0:.*]] = arith.constant 0 : index
+// CHECK:         %[[C1:.*]] = arith.constant 1 : index
+// CHECK:         %[[C2:.*]] = arith.constant 2 : index
+// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C0]], %[[C0]]]
+// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C0]], %[[C1]]]
+// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C1]], %[[C0]]]
+// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C1]], %[[C1]]]
+// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C2]], %[[C0]]]
+// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C2]], %[[C1]]]
+// CHECK:         %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
+// CHECK:         return %[[RET]] : tensor<3x2xindex>
+func @tensor.from_elements_2d(%arg0: index, %arg1: index) -> tensor<3x2xindex> {
+  %0 = tensor.from_elements %arg0, %arg1, %arg0, %arg1, %arg0, %arg1
+         : tensor<3x2xindex>
+  return %0 : tensor<3x2xindex>
+}
+
+// CHECK-LABEL: func @tensor.from_elements_3d(
+// CHECK-SAME:      %[[ELEM0:.*]]: f32, %[[ELEM1:.*]]: f32, %[[ELEM2:.*]]: f32)
+// CHECK-SAME:      -> tensor<3x2x2xf32> {
+// CHECK:         %[[MEMREF:.*]] = memref.alloc() : memref<3x2x2xf32>
+// CHECK:         %[[C0:.*]] = arith.constant 0 : index
+// CHECK:         %[[C1:.*]] = arith.constant 1 : index
+// CHECK:         %[[C2:.*]] = arith.constant 2 : index
+// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C0]], %[[C0]], %[[C0]]]
+// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C0]], %[[C0]], %[[C1]]]
+// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C0]], %[[C1]], %[[C0]]]
+// CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C0]], %[[C1]], %[[C1]]]
+// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C1]], %[[C0]], %[[C0]]]
+// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C1]], %[[C0]], %[[C1]]]
+// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C1]], %[[C1]], %[[C0]]]
+// CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C1]], %[[C1]], %[[C1]]]
+// CHECK:         store %[[ELEM2]], %[[MEMREF]][%[[C2]], %[[C0]], %[[C0]]]
+// CHECK:         store %[[ELEM2]], %[[MEMREF]][%[[C2]], %[[C0]], %[[C1]]]
+// CHECK:         store %[[ELEM2]], %[[MEMREF]][%[[C2]], %[[C1]], %[[C0]]]
+// CHECK:         store %[[ELEM2]], %[[MEMREF]][%[[C2]], %[[C1]], %[[C1]]]
+// CHECK:         %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
+// CHECK:         return %[[RET]] : tensor<3x2x2xf32>
+func @tensor.from_elements_3d(%f0: f32, %f1: f32, %f2: f32)
+    -> tensor<3x2x2xf32> {
+  %0 = tensor.from_elements %f0,%f0,%f0,%f0,%f1,%f1,%f1,%f1,%f2,%f2,%f2,%f2
+         : tensor<3x2x2xf32>
+  return %0 : tensor<3x2x2xf32>
+}
+
 // CHECK-LABEL:   func @tensor.generate(
 // CHECK-SAME:                                       %[[ARG:.*]]: tensor<*xf32>,
 // CHECK-SAME:                                       %[[DYNAMIC_EXTENT:.*]]: index) -> tensor<?xindex> {
 // CHECK:           %[[CASTED:.*]] = bufferization.to_memref %[[ARG]] : memref<*xf32>
 // CHECK:           %[[MEMREF:.*]] = memref.alloc(%[[DYNAMIC_EXTENT]]) : memref<?xindex>
 // CHECK:           %[[C0:.*]] = arith.constant 0 : index
 // CHECK:           %[[C1:.*]] = arith.constant 1 : index
 // CHECK:           scf.parallel (%[[I:.*]]) = (%[[C0]]) to (%[[DYNAMIC_EXTENT]]) step (%[[C1]]) {

mlir/test/Dialect/Tensor/canonicalize.mlir

func @extract_from_tensor.from_elements(%element : index) -> index {
   %tensor = tensor.from_elements %element : tensor<1xindex>
   %extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex>
   // CHECK: [[ARG]] : index
   return %extracted_element : index
 }
 
 // -----
 
+// CHECK-LABEL: func @extract_from_tensor.from_elements_2d
+func @extract_from_tensor.from_elements_2d(%el1 : index, %el2 : index)
+    -> index {
+  // CHECK-SAME: ([[EL1:%.*]]: index, [[EL2:%.*]]: index)
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %tensor = tensor.from_elements %el1, %el1, %el2, %el2 : tensor<2x2xindex>
+  %extracted_element = tensor.extract %tensor[%c1, %c0] : tensor<2x2xindex>
+  // CHECK: [[EL2]] : index
+  return %extracted_element : index
+}
+
+// -----
+
 // Ensure the optimization doesn't segfault from bad constants
 // CHECK-LABEL: func @extract_negative_from_tensor.from_elements
 func @extract_negative_from_tensor.from_elements(%element : index) -> index {
   // CHECK-SAME: ([[ARG:%.*]]: index)
   %c-1 = arith.constant -1 : index
   %tensor = tensor.from_elements %element : tensor<1xindex>
   %extracted_element = tensor.extract %tensor[%c-1] : tensor<1xindex>
   // CHECK: tensor.from_elements

mlir/test/Dialect/Tensor/invalid.mlir

func @insert_too_many_indices(%arg0: f32, %arg1: tensor<?xf32>) {
   // expected-error@+1 {{incorrect number of indices}}
   %0 = tensor.insert %arg0 into %arg1[] : tensor<?xf32>
   return
 }
 
 // -----
 
 func @tensor.from_elements_wrong_result_type() {
-  // expected-error@+2 {{'result' must be 1D tensor of any type values, but got 'tensor<*xi32>'}}
+  // expected-error@+2 {{'result' must be statically shaped tensor of any type values, but got 'tensor<*xi32>'}}
   %c0 = arith.constant 0 : i32
   %0 = tensor.from_elements %c0 : tensor<*xi32>
   return
 }
 
 // -----
 
 func @tensor.from_elements_wrong_elements_count() {

mlir/test/Dialect/Tensor/ops.mlir

func @insert(%arg0: f32, %arg1: index, %arg2: tensor<?x?x?xf32>, %arg3: tensor<*xf32>) {
   // CHECK: tensor.insert %[[SCALAR]] into %[[DEST2]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<*xf32>
   %1 = tensor.insert %arg0 into %arg3[%arg1, %arg1, %arg1] : tensor<*xf32>
   return
 }
 
 // CHECK-LABEL: func @tensor.from_elements() {
 func @tensor.from_elements() {
   %c0 = "arith.constant"() {value = 0: index} : () -> index
-  // CHECK: %0 = tensor.from_elements %c0 : tensor<1xindex>
+  // CHECK: tensor.from_elements %c0 : tensor<1xindex>
   %0 = tensor.from_elements %c0 : tensor<1xindex>
 
   %c1 = "arith.constant"() {value = 1: index} : () -> index
-  // CHECK: %1 = tensor.from_elements %c0, %c1 : tensor<2xindex>
+  // CHECK: tensor.from_elements %c0, %c1 : tensor<2xindex>
   %1 = tensor.from_elements %c0, %c1 : tensor<2xindex>
 
   %c0_f32 = "arith.constant"() {value = 0.0: f32} : () -> f32
   // CHECK: [[C0_F32:%.*]] = arith.constant
-  // CHECK: %2 = tensor.from_elements [[C0_F32]] : tensor<1xf32>
+  // CHECK: 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>
 
+  // CHECK: tensor.from_elements %c0, %c1, %c0, %c1, %c0, %c1 : tensor<2x3xindex>
+  %4 = tensor.from_elements %c0, %c1, %c0, %c1, %c0, %c1 : tensor<2x3xindex>
+
+  // CHECK: tensor.from_elements %c0 : tensor<index>
+  %5 = tensor.from_elements %c0 : tensor<index>
   return
 }
 
 // CHECK-LABEL: @tensor.generate
 func @tensor.generate(%m : index, %n : index)
     -> tensor<?x3x?xf32> {
   %tnsr = tensor.generate %m, %n {
     ^bb0(%i : index, %j : index, %k : index):

mlir/test/lib/Dialect/Test/TestDialect.cpp

 }
 
 LogicalResult OpWithResultShapeInterfaceOp::reifyReturnTypeShapes(
     OpBuilder &builder, ValueRange operands,
     llvm::SmallVectorImpl<Value> &shapes) {
   Location loc = getLoc();
   shapes.reserve(operands.size());
   for (Value operand : llvm::reverse(operands)) {
-    auto currShape = llvm::to_vector<4>(llvm::map_range(
-        llvm::seq<int64_t>(
-            0, operand.getType().cast<RankedTensorType>().getRank()),
-        [&](int64_t dim) -> Value {
+    auto rank = operand.getType().cast<RankedTensorType>().getRank();
+    auto currShape = llvm::to_vector<4>(
+        llvm::map_range(llvm::seq<int64_t>(0, rank), [&](int64_t dim) -> Value {
           return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
         }));
     shapes.push_back(builder.create<tensor::FromElementsOp>(
-        getLoc(), builder.getIndexType(), currShape));
+        getLoc(), RankedTensorType::get({rank}, builder.getIndexType()),
+        currShape));
   }
   return success();
 }
 
 LogicalResult OpWithResultShapePerDimInterfaceOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
   Location loc = getLoc();
   shapes.reserve(getNumOperands());