[MLIR][Standard] Add `dynamic_tensor_from_elements` operation

[frgossen]

Summary

With dynamic_tensor_from_elements tensor values of dynamic size can be
created. The body of the operation essentially maps the index space to tensor
elements.

Declare SCF operations in the scf namespace to avoid name clash with the new
std.yield operation. Resolve ambiguities between linalg/shape/std/scf.yield
operations.

[silvas]

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

1481

We need to verify that all ops in the body are NoSideEffect

1482

I think this needs some more traits related to regions (look at scf dialect ops)

1504

we should use the term "extent" instead of "dimension" in this case, as "dimension" is ambiguous (for example, the meaning of "dimension" in a transpose op)

3262

Needs HasParent, NoSideEffect (see scf dialect)

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

1666

attr-dict handling is needed?

1677

personally, I have found this "only take the dynamic dimensions" to be a huge pain in other ops (requires lots of special handling code). I would prefer to take all dimensions.

As an example, if we have tensor_cast %result_of_dynamic_tensor_from_elements : tensor<?xf32> to tensor<4xf32>, we cannot elide the tensor_cast without major repair of the dynamic_tensor_from_elements op.

We can have a canonicalization that looks at the output type, and replaces any input extents with constants if the output type is more precise.

1684

reword to something like "all body operands must be index" + "must have one body argument per input dimension"

mlir/test/Dialect/Standard/invalid.mlir

62

this case can be verified with traits

[silvas]

Also, can you update the commit description to include a description of the namespace changes?

[frgossen]

Address comments

[frgossen]

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

1666

Not at this point.

Correct me if I'm wrong but I thought attr-dict should always be supported. At least assemblyFormat enforces this.

1677

Why would this be a major repair?
With the current version, you'd have to erase one operand and otherwise you'd have to replace it with a constant. The body arguments span the entire index space anyways.
I don't yet see the advantage and think it'd be nice if this was consistent with, e.g. alloc.

[frgossen]

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

1677

Just a follow up comment:
When implementing the lowering for this op in TensorFlow, I found it very useful that that alloca expect the same operands.

[herhut]

Thanks. Please fix comments before landing.

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

3272

Would let verifier = ? work to avoid specifying the empty hook?

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

1677

I also think we should strive for consistency here. If we want to change this kind of behavior, we need to do it globally and not just for this op.

1684

I would also prefer distinct error messages for the two failure cases (too few arguments and wrong type).

3315

See my other comment. This should not be needed.

[frgossen]

Address comments

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

1684

Sorry, misunderstood that comment.

mlir/include/mlir/Dialect/SCF/SCFOps.td

 #define MLIR_DIALECT_SCF_SCFOPS
 
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def SCF_Dialect : Dialect {
   let name = "scf";
-  let cppNamespace = "";
+  let cppNamespace = "scf";
 }
 
 // Base class for SCF dialect ops.
 class SCF_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<SCF_Dialect, mnemonic, traits> {
   // For every standard op, there needs to be a:
   //   * void print(OpAsmPrinter &p, ${C++ class of Op} op)
   //   * LogicalResult verify(${C++ class of Op} op)
   //   * ParseResult parse${C++ class of Op}(OpAsmParser &parser,
   //                                         OperationState &result)
   // functions.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 def ForOp : SCF_Op<"for",
       [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-       SingleBlockImplicitTerminator<"YieldOp">,
+       SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
   let description = [{
     The "scf.for" operation represents a loop taking 3 SSA value as operands
     that represent the lower bound, upper bound and step respectively.  The
     operation defines an SSA value for its induction variable. It has one
     region capturing the loop body. The induction variable is represented as an
     argument of this region. This SSA value always has type index, which is the
let extraClassDeclaration = [{
     /// induction variable. LoopOp only has one region, so 0 is the only valid
     /// value for `index`.
     OperandRange getSuccessorEntryOperands(unsigned index);
   }];
 }
 
 def IfOp : SCF_Op<"if",
       [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-       SingleBlockImplicitTerminator<"YieldOp">, RecursiveSideEffects,
+       SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveSideEffects,
        NoRegionArguments]> {
   let summary = "if-then-else operation";
   let description = [{
     The `scf.if` operation represents an if-then-else construct for
     conditionally executing two regions of code. The operand to an if operation
     is a boolean value. For example:
 
     ```mlir
let extraClassDeclaration = [{
     }
   }];
 }
 
 def ParallelOp : SCF_Op<"parallel",
     [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      RecursiveSideEffects,
-     SingleBlockImplicitTerminator<"YieldOp">]> {
+     SingleBlockImplicitTerminator<"scf::YieldOp">]> {
   let summary = "parallel for operation";
   let description = [{
     The "scf.parallel" operation represents a loop nest taking 4 groups of SSA
     values as operands that represent the lower bounds, upper bounds, steps and
     initial values, respectively. The operation defines a variadic number of
     SSA values for its induction variables. It has one region capturing the
     loop body. The induction variables are represented as an argument of this
     region. These SSA values always have type index, which is the size of the

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

 // DivFOp
 //===----------------------------------------------------------------------===//
 
 def DivFOp : FloatArithmeticOp<"divf"> {
   let summary = "floating point division operation";
 }
 
 //===----------------------------------------------------------------------===//
+// DynamicTensorFromElementsOp
+//===----------------------------------------------------------------------===//
+
+def DynamicTensorFromElementsOp : Std_Op<"dynamic_tensor_from_elements",
+    [NoSideEffect]> {
+  string summary = "Creates a dynamically sized tensor from elements";
+  string description = [{
+    This operation creates a dynamically sized tensor with elements of any type.
+    It expects one index operand per dynamic dimension of the result tensor,
+    each defining the corresponding extent at runtime.
+
+    The body region defines the tensor's elements. It takes index operands as
+    its region arguments that span the index space. The element at the given
+    position is yielded with the `yield` operation (see `YieldOp`).
+
+    Example:
+
+    ```mlir
+      %tnsr = dynamic_tensor_from_elements %m, %n {
+      ^bb0(%i : index, %j : index, %k : index):
+        ...
+        yield %elem : f32
+      } : tensor<?x3x?f32>
+    ```
+  }];
+
+  let arguments = (ins Variadic<Index>:$dynamicDimensions);
+  let results = (outs AnyRankedTensor:$result);
+  let regions = (region SizedRegion<1>:$body);
+}
+
+//===----------------------------------------------------------------------===//
 // ExpOp
 //===----------------------------------------------------------------------===//
 
 def ExpOp : FloatUnaryOp<"exp"> {
   let summary = "base-e exponential of the specified value";
   let description = [{
     Syntax:
 
operand_range getDynamicSizes() {
       return {sizes().begin(), sizes().end()};
     }
   }];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// YieldOp
+//===----------------------------------------------------------------------===//
+
+def YieldOp : Std_Op<"yield", [NoSideEffect, Terminator]> {
+  let summary = "Yield a value from a region";
+  let description = [{
+     This operation is used to yield a single value from a within a region. It
+     is used to create dynamically sized tensors
+     (see `DynamicTensorFromElementsOp`).
+  }];
+
+  let arguments = (ins AnyType:$value);
+  let assemblyFormat = "$value attr-dict `:` type($value)";
+}
+
+//===----------------------------------------------------------------------===//
 // XOrOp
 //===----------------------------------------------------------------------===//
 
 def XOrOp : IntArithmeticOp<"xor", [Commutative]> {
   let summary = "integer binary xor";
   let description = [{
     The `xor` operation takes two operands and returns one result, each of these
     is required to be the same type. This type may be an integer scalar type, a

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     return success();
   }
 };
 
 // YieldOp produces and LLVM::ReturnOp.
 class YieldOpConversion : public ConvertToLLVMPattern {
 public:
   explicit YieldOpConversion(MLIRContext *context, LLVMTypeConverter &lowering_)
-      : ConvertToLLVMPattern(YieldOp::getOperationName(), context, lowering_) {}
+      : ConvertToLLVMPattern(linalg::YieldOp::getOperationName(), context,
+                             lowering_) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, operands);
     return success();
   }
 };

mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp

if (first) {
       // Store the results of the outermost loop that will be used to replace
       // the results of the parallel loop when it is fully rewritten.
       loopResults.assign(forOp.result_begin(), forOp.result_end());
       first = false;
     } else if (!forOp.getResults().empty()) {
       // A loop is constructed with an empty "yield" terminator if there are
       // no results.
       rewriter.setInsertionPointToEnd(rewriter.getInsertionBlock());
-      rewriter.create<YieldOp>(loc, forOp.getResults());
+      rewriter.create<scf::YieldOp>(loc, forOp.getResults());
     }
 
     rewriter.setInsertionPointToStart(forOp.getBody());
   }
 
   // Now copy over the contents of the body.
   SmallVector<Value, 4> yieldOperands;
   yieldOperands.reserve(parallelOp.getNumResults());
for (auto &op : parallelOp.getBody()->without_terminator()) {
     for (auto &nested : reduceBlock.without_terminator())
       rewriter.clone(nested, mapping);
     yieldOperands.push_back(
         mapping.lookup(reduceBlock.getTerminator()->getOperand(0)));
   }
 
   if (!yieldOperands.empty()) {
     rewriter.setInsertionPointToEnd(rewriter.getInsertionBlock());
-    rewriter.create<YieldOp>(loc, yieldOperands);
+    rewriter.create<scf::YieldOp>(loc, yieldOperands);
   }
 
   rewriter.replaceOp(parallelOp, loopResults);
 
   return success();
 }
 
 void mlir::populateLoopToStdConversionPatterns(

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

result.addAttribute(TransposeOp::getPermutationAttrName(),
                       AffineMapAttr::get(permutation));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, YieldOp op) {
+static void print(OpAsmPrinter &p, linalg::YieldOp op) {
   p << op.getOperationName();
   if (op.getNumOperands() > 0)
     p << ' ' << op.getOperands();
   p.printOptionalAttrDict(op.getAttrs());
   if (op.getNumOperands() > 0)
     p << " : " << op.getOperandTypes();
 }
 
 static ParseResult parseYieldOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 2> opInfo;
   SmallVector<Type, 2> types;
   llvm::SMLoc loc = parser.getCurrentLocation();
   return failure(parser.parseOperandList(opInfo) ||
                  parser.parseOptionalAttrDict(result.attributes) ||
                  (!opInfo.empty() && parser.parseColonTypeList(types)) ||
                  parser.resolveOperands(opInfo, types, loc, result.operands));
 }
 
 // Check the operand number and types must match the element types of the
 // LinalgOp interface's shaped operands.
-static LogicalResult verifyYield(YieldOp op, LinalgOp linalgOpInterface) {
+static LogicalResult verifyYield(linalg::YieldOp op,
+                                 LinalgOp linalgOpInterface) {
   auto nOutputs = linalgOpInterface.getNumOutputs();
   if (op.getNumOperands() != nOutputs)
     return op.emitOpError("expected number of yield values (")
            << nOutputs << ") to match the number of operands of the enclosing "
            << "LinalgOp (" << op.getNumOperands() << ")";
 
   for (unsigned i = 0; i != nOutputs; ++i) {
     auto elementType =
         linalgOpInterface.getOutputShapedType(i).getElementType();
     if (op.getOperand(i).getType() != elementType)
       return op.emitOpError("type of yield operand ")
              << (i + 1) << " (" << op.getOperand(i).getType()
              << ") doesn't match "
              << "the element type of the enclosing linalg.generic op ("
              << elementType << ")";
   }
   return success();
 }
 
-static LogicalResult verify(YieldOp op) {
+static LogicalResult verify(linalg::YieldOp op) {
   auto *parentOp = op.getParentOp();
   if (parentOp->getNumRegions() != 1 || parentOp->getRegion(0).empty())
     return op.emitOpError("expected single non-empty parent region");
 
   if (auto linalgOp = dyn_cast<LinalgOp>(parentOp))
     return verifyYield(op, cast<LinalgOp>(parentOp));
 
   return op.emitOpError("expected parent op with LinalgOp interface");

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

for (auto consumerArg : llvm::enumerate(consumerBlock.getArguments())) {
         mapper.map(consumerArg.value(),
                    fusedBlock->addArgument(consumerArg.value().getType()));
       }
     }
 
     // Add operations from producer (except the yield operation) to the fused
     // op.
     for (auto &op : producerBlock.getOperations()) {
-      if (auto yieldOp = dyn_cast<YieldOp>(op)) {
+      if (auto yieldOp = dyn_cast<linalg::YieldOp>(op)) {
         // Lookup the value the yield operation is mapped to.
         Value yieldVal = yieldOp.getOperand(0);
         if (Value clonedVal = mapper.lookupOrNull(yieldVal))
           mapper.map(
               consumerBlock.getArgument(consumerIdx + numConsumerIndices),
               clonedVal);
         continue;
       }

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

static void inlineRegionAndEmitStore(OpType op, ArrayRef<Value> indexedValues,
   map.map(block.getArguments(), indexedValues);
   for (auto &op : block.without_terminator()) {
     assert(op.getNumRegions() == 0 && "expected a non-nested region");
     auto *newOp = b.clone(op, map);
     map.map(op.getResults(), newOp->getResults());
   }
 
   Operation &terminator = block.back();
-  assert(isa<YieldOp>(terminator) &&
+  assert(isa<linalg::YieldOp>(terminator) &&
          "expected a yield op in the end of the region");
   for (unsigned i = 0, e = terminator.getNumOperands(); i < e; ++i) {
     IndexedValueType O(outputBuffers[i]);
     O(indexing[i]) = map.lookupOrDefault(terminator.getOperand(i));
   }
 }
 
 // Returns a pair that contains input indices and output indices of a

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

if (!llvm::hasNItems(r.front().begin(), r.front().end(), 3))
     return false;
 
   using mlir::matchers::m_Val;
   auto a = m_Val(r.getArgument(0));
   auto b = m_Val(r.getArgument(1));
   auto c = m_Val(r.getArgument(2));
   // TODO: Update this detection once we have  matcher support for specifying
   // that any permutation of operands matches.
-  auto pattern1 = m_Op<YieldOp>(m_Op<AddFOp>(m_Op<MulFOp>(a, b), c));
-  auto pattern2 = m_Op<YieldOp>(m_Op<AddFOp>(c, m_Op<MulFOp>(a, b)));
-  auto pattern3 = m_Op<YieldOp>(m_Op<AddFOp>(m_Op<MulFOp>(b, a), c));
-  auto pattern4 = m_Op<YieldOp>(m_Op<AddFOp>(c, m_Op<MulFOp>(b, a)));
-  auto pattern5 = m_Op<YieldOp>(m_Op<AddIOp>(m_Op<MulIOp>(a, b), c));
-  auto pattern6 = m_Op<YieldOp>(m_Op<AddIOp>(c, m_Op<MulIOp>(a, b)));
-  auto pattern7 = m_Op<YieldOp>(m_Op<AddIOp>(m_Op<MulIOp>(b, a), c));
-  auto pattern8 = m_Op<YieldOp>(m_Op<AddIOp>(c, m_Op<MulIOp>(b, a)));
+  auto pattern1 = m_Op<linalg::YieldOp>(m_Op<AddFOp>(m_Op<MulFOp>(a, b), c));
+  auto pattern2 = m_Op<linalg::YieldOp>(m_Op<AddFOp>(c, m_Op<MulFOp>(a, b)));
+  auto pattern3 = m_Op<linalg::YieldOp>(m_Op<AddFOp>(m_Op<MulFOp>(b, a), c));
+  auto pattern4 = m_Op<linalg::YieldOp>(m_Op<AddFOp>(c, m_Op<MulFOp>(b, a)));
+  auto pattern5 = m_Op<linalg::YieldOp>(m_Op<AddIOp>(m_Op<MulIOp>(a, b), c));
+  auto pattern6 = m_Op<linalg::YieldOp>(m_Op<AddIOp>(c, m_Op<MulIOp>(a, b)));
+  auto pattern7 = m_Op<linalg::YieldOp>(m_Op<AddIOp>(m_Op<MulIOp>(b, a), c));
+  auto pattern8 = m_Op<linalg::YieldOp>(m_Op<AddIOp>(c, m_Op<MulIOp>(b, a)));
   return pattern1.match(&r.front().back()) ||
          pattern2.match(&r.front().back()) ||
          pattern3.match(&r.front().back()) ||
          pattern4.match(&r.front().back()) ||
          pattern5.match(&r.front().back()) ||
          pattern6.match(&r.front().back()) ||
          pattern7.match(&r.front().back()) || pattern8.match(&r.front().back());
 }

mlir/lib/Dialect/SCF/SCF.cpp

struct SCFInlinerInterface : public DialectInlinerInterface {
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }
   // Handle the given inlined terminator by replacing it with a new operation
   // as necessary. Required when the region has only one block.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {
-    auto retValOp = dyn_cast<YieldOp>(op);
+    auto retValOp = dyn_cast<scf::YieldOp>(op);
     if (!retValOp)
       return;
 
     for (auto retValue : llvm::zip(valuesToRepl, retValOp.getOperands())) {
       std::get<0>(retValue).replaceAllUsesWith(std::get<1>(retValue));
     }
   }
 };
static ParseResult parseYieldOp(OpAsmParser &parser, OperationState &result) {
   // values.
   if (parser.parseOperandList(operands) ||
       parser.parseOptionalColonTypeList(types) ||
       parser.resolveOperands(operands, types, loc, result.operands))
     return failure();
   return success();
 }
 
-static void print(OpAsmPrinter &p, YieldOp op) {
+static void print(OpAsmPrinter &p, scf::YieldOp op) {
   p << op.getOperationName();
   if (op.getNumOperands() != 0)
     p << ' ' << op.getOperands() << " : " << op.getOperandTypes();
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
+namespace mlir {
+namespace scf {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/SCFOps.cpp.inc"
+} // namespace scf
+} // namespace mlir

mlir/lib/Dialect/Shape/IR/Shape.cpp

void SizeToIndexOp::getCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<IndexToSizeToIndexCanonicalization>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(YieldOp op) {
+static LogicalResult verify(shape::YieldOp op) {
   auto *parentOp = op.getParentOp();
   auto results = parentOp->getResults();
   auto operands = op.getOperands();
 
   if (parentOp->getNumResults() != op.getNumOperands())
     return op.emitOpError() << "number of operands does not match number of "
                                "results of its parent";
   for (auto e : llvm::zip(results, operands))

mlir/lib/Dialect/Shape/Transforms/ShapeToShapeLowering.cpp

Value init = op.getDialect()
                    ->getResult(0);
   ReduceOp reduce = rewriter.create<ReduceOp>(loc, op.shape(), init);
 
   // Generate reduce operator.
   Block *body = reduce.getBody();
   OpBuilder b = OpBuilder::atBlockEnd(body);
   Value product = b.create<MulOp>(loc, valueType, body->getArgument(1),
                                   body->getArgument(2));
-  b.create<YieldOp>(loc, product);
+  b.create<shape::YieldOp>(loc, product);
 
   rewriter.replaceOp(op, reduce.result());
   return success();
 }
 
 namespace {
 struct ShapeToShapeLowering
     : public ShapeToShapeLoweringBase<ShapeToShapeLowering> {

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

 
 Optional<int64_t> DimOp::getConstantIndex() {
   if (auto constantOp = index().getDefiningOp<ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 
 static LogicalResult verify(DimOp op) {
-
   // Assume unknown index to be in range.
   Optional<int64_t> index = op.getConstantIndex();
   if (!index.hasValue())
     return success();
 
   // Check that constant index is not knowingly out of range.
   auto type = op.memrefOrTensor().getType();
   if (auto tensorType = type.dyn_cast<RankedTensorType>()) {
LogicalResult DmaWaitOp::verify() {
   if (!getNumElements().getType().isIndex())
     return emitOpError()
            << "expected the number of elements to be of index type";
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// DynamicTensorFromElementsOp
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseDynamicTensorFromElementsOp(OpAsmParser &parser,
+                                                    OperationState &result) {
+  // Parse operands.
+  SmallVector<OpAsmParser::OperandType, 4> dynamicDimensions;
+  Type indexTy = parser.getBuilder().getIndexType();
+  if (parser.parseOperandList(dynamicDimensions) ||
+      parser.resolveOperands(dynamicDimensions, indexTy, result.operands))
+    return failure();
+
+  // Parse body.
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, {}, {}))
+    return failure();
+
+  // Parse result type.
+  Type resultType;
+  if (parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseColonType(resultType))
+    return failure();
+  result.addTypes(resultType);
+
+  return success();
+}
+
+static void print(OpAsmPrinter &p, DynamicTensorFromElementsOp op) {
+  p << "dynamic_tensor_from_elements " << op.dynamicDimensions();
+  p.printRegion(op.body());
+  p.printOptionalAttrDict(op.getAttrs());
+  p << " : " << op.getType();
+}
+
+static LogicalResult verify(DynamicTensorFromElementsOp op) {
+  // Ensure that the tensor type has as many dynamic dimensions as are specified
+  // by the operands.
+  RankedTensorType resultTy = op.getType().cast<RankedTensorType>();
+  if (op.getNumOperands() != resultTy.getNumDynamicDims())
+    return op.emitError(
+        "must have as many index operands as dynamic dimensions in the result "
+        "type");
+
+  // Ensure that region arguments span the index space.
+  if (op.body().getNumArguments() != resultTy.getRank() ||
+      !llvm::all_of(op.body().getArgumentTypes(),
+                    [](Type ty) { return ty.isIndex(); }))
+    return op.emitError("body arguments must span index space");
+
+  // Ensure that the region yields an element of the right type.
+  auto yieldOp =
+      llvm::dyn_cast<YieldOp>(op.body().getBlocks().front().getTerminator());
+  if (!yieldOp || yieldOp.value().getType() != resultTy.getElementType())
+    return op.emitOpError(
+        "body must be terminated with a `yield` operation of the tensor "
+        "element type");
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ExtractElementOp op) {
   // Verify the # indices match if we have a ranked type.
   auto aggregateType = op.getAggregate().getType().cast<ShapedType>();
   if (aggregateType.hasRank() &&
       aggregateType.getRank() != op.getNumOperands() - 1)
OpFoldResult XOrOp::fold(ArrayRef<Attribute> operands) {
   if (lhs() == rhs())
     return Builder(getContext()).getZeroAttr(getType());
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a ^ b; });
 }
 
 //===----------------------------------------------------------------------===//
+// YieldOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(YieldOp op) { return success(); }
+
+//===----------------------------------------------------------------------===//
 // ZeroExtendIOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ZeroExtendIOp op) {
   auto srcType = getElementTypeOrSelf(op.getOperand().getType());
   auto dstType = getElementTypeOrSelf(op.getType());
 
   if (srcType.isa<IndexType>())

mlir/test/Dialect/Standard/invalid.mlir

 // -----
 
 // CHECK-LABEL: test_index_cast_tensor_error
 func @test_index_cast_tensor_error(%arg0 : tensor<index>) -> i64 {
   // expected-error @+1 {{requires the same shape for all operands and results}}
   %0 = index_cast %arg0 : tensor<index> to i64
   return %0 : i64
 }
+
+// -----
+
+func @dynamic_tensor_from_elements(%m : index)
+    -> tensor<?x3x?xf32> {
+  // expected-error @+1 {{must have as many index operands as dynamic dimensions in the result type}}
+  %tnsr = dynamic_tensor_from_elements %m {
+    ^bb0(%i : index, %j : index, %k : index):
+      %elem = constant 8.0 : f32
+      yield %elem : f32
+  } : tensor<?x3x?xf32>
+  return %tnsr : tensor<?x3x?xf32>
+}
+
+// -----
+
+func @dynamic_tensor_from_elements(%m : index, %n : index)
+    -> tensor<?x3x?xf32> {
+  // expected-error @+1 {{body arguments must span index space}}
+  %tnsr = dynamic_tensor_from_elements %m, %n {
+    ^bb0(%i : index, %j : index):
+      %elem = constant 8.0 : f32
+      yield %elem : f32
+  } : tensor<?x3x?xf32>
+  return %tnsr : tensor<?x3x?xf32>
+}
+
+// -----
+
+func @dynamic_tensor_from_elements(%m : index, %n : index)
+    -> tensor<?x3x?xf32> {
+  // expected-error @+1 {{body arguments must span index space}}
+  %tnsr = dynamic_tensor_from_elements %m, %n {
+    ^bb0(%i : index, %j : index, %k : i64):
+      %elem = constant 8.0 : f32
+      yield %elem : f32
+  } : tensor<?x3x?xf32>
+  return %tnsr : tensor<?x3x?xf32>
+}
+
+// -----
+
+func @dynamic_tensor_from_elements(%m : index, %n : index)
+    -> tensor<?x3x?xf32> {
+  // expected-error @+1 {{body must be terminated with a `yield` operation of the tensor element type}}
+  %tnsr = dynamic_tensor_from_elements %m, %n {
+    ^bb0(%i : index, %j : index, %k : index):
+      %elem = constant 8.0 : f32
+      return %elem : f32
+  } : tensor<?x3x?xf32>
+  return %tnsr : tensor<?x3x?xf32>
+}
+
+// -----
+
+func @dynamic_tensor_from_elements(%m : index, %n : index)
+    -> tensor<?x3x?xf32> {
+  // expected-error @+1 {{body must be terminated with a `yield` operation of the tensor element type}}
+  %tnsr = dynamic_tensor_from_elements %m, %n {
+    ^bb0(%i : index, %j : index, %k : index):
+      %elem = constant 8 : i32
+      yield %elem : i32
+  } : tensor<?x3x?xf32>
+  return %tnsr : tensor<?x3x?xf32>
+}

mlir/test/Dialect/Standard/ops.mlir

-// RUN: mlir-opt -split-input-file %s | FileCheck %s
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: test_index_cast
 func @test_index_cast(%arg0 : index) -> i64 {
   %0 = index_cast %arg0 : index to i64
   return %0 : i64
 }
 
 // CHECK-LABEL: test_index_cast_tensor
 func @test_index_cast_tensor(%arg0 : tensor<index>) -> tensor<i64> {
   %0 = index_cast %arg0 : tensor<index> to tensor<i64>
   return %0 : tensor<i64>
 }
 
 // CHECK-LABEL: test_index_cast_tensor_reverse
 func @test_index_cast_tensor_reverse(%arg0 : tensor<i64>) -> tensor<index> {
   %0 = index_cast %arg0 : tensor<i64> to tensor<index>
   return %0 : tensor<index>
 }
 
 func @assert(%arg : i1) {
   assert %arg, "Some message in case this assertion fails."
   return
 }
+
+func @dynamic_tensor_from_elements(%m : index, %n : index)
+    -> tensor<?x3x?xf32> {
+  %tnsr = dynamic_tensor_from_elements %m, %n {
+    ^bb0(%i : index, %j : index, %k : index):
+      %elem = constant 8.0 : f32
+      yield %elem : f32
+  } : tensor<?x3x?xf32>
+  return %tnsr : tensor<?x3x?xf32>
+}
+