Introduce tensor.insert op to Tensor dialect.

[hanchung]

Summary

Add tensor.insert op to make tensor.extract/tensor.insert work in pairs
for scalar domain. Like subtensor/subtensor_insert work in pairs in
tensor domain, and vector.transfer_read/vector.transfer_write work in
pairs in vector domain.

[hanchung]

delete wrong comments

[silvas]

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

212

Could the syntax tensor.insert %ut into %dest[%1, %2] : tensor<*xi32>work as well? The scalar type being inserted should always be the tensor element type (I think TypesMatchWith is able to support this)

227

nit: why is cast to ShapedType needed?

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

299

Redundant after you add TypesMatchWith for scalar.

[hanchung]

address comments

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

212

Yes, it works. I also added the case to ops.mlir. It won't complain after adding TypesMatchWith.

227

Done. It is not needed. I copied it from ExtractOp, but did not notice it.

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

299

Done, deleted.

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

let builders = [
     OpBuilder<(ins "Type":$resultTy, "ValueRange":$dynamicExtents,
       "function_ref<void(OpBuilder &, Location, ValueRange)>")>,
   ];
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// InsertOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_InsertOp : Tensor_Op<"insert",
+    [NoSideEffect,
+     TypesMatchWith<"result type matches type of dest",
+                    "dest", "result",
+                    "$_self.cast<ShapedType>()">,
+     TypesMatchWith<"scalar type matches element type of dest",
+                    "dest", "scalar",
+                    "$_self.cast<ShapedType>().getElementType()">]> {
+  let summary = "element insertion operation";
+  let description = [{
+    The `tensor.insert` op writes a tensor into a tensor `dest`as specified by
+    the operation's indices.
+
+    It returns a copy of `dest` with the proper subtensor updated with the value
+    of `scalar`.
+
+    The arity of indices must match the rank of the tensor `dest` (i.e., if a
+    tensor is of rank 3, then 3 indices are required for the extract. The
+    indices should all be of `index` type.
+
+    Example:
+
+    ```mlir
+    %4 = tensor.insert %t into %dest[%1, %2] : tensor<4x4xi32>
+    %5 = tensor.insert %rt into %dest[%1, %2] : tensor<?x?xi32>
+    %6 = tensor.insert %ut into %dest[%1, %2] : tensor<*xi32>
+    ```
+  }];
+
+  let arguments = (ins AnyType:$scalar,
+                       AnyTensor:$dest,
+                       Variadic<Index>:$indices);
+  let results = (outs AnyTensor:$result);
+  let assemblyFormat = [{
+    $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$scalar, "Value":$dest,
+      CArg<"ValueRange", "{}">:$indices), [{
+      auto resType = dest.getType();
+      build($_builder, $_state, resType, scalar, dest, indices);
+    }]>];
+
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_ReshapeOp: Tensor_Op<"reshape", [NoSideEffect]>  {
   let summary = "tensor reshape operation";
   let description = [{
     The `reshape` operation converts a tensor from one type to an equivalent
     type with a provided shape. The source and destination types are compatible

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

 } // namespace
 
 void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {
   results.add<ExtractElementFromTensorFromElements>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// InsertOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(InsertOp op) {
+  // Verify the # indices match if we have a ranked type.
+  if (auto destType = op.dest().getType().dyn_cast<RankedTensorType>())
+    if (destType.getRank() != static_cast<int64_t>(op.indices().size()))
+      return op.emitOpError("incorrect number of indices");
+  return success();
+}
+
+OpFoldResult InsertOp::fold(ArrayRef<Attribute> operands) {
+  Attribute scalar = operands[0];
+  Attribute dest = operands[1];
+  if (scalar && dest)
+    if (auto splatDest = dest.dyn_cast<SplatElementsAttr>())
+      if (scalar == splatDest.getSplatValue())
+        return dest;
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // GenerateOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(GenerateOp 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())

mlir/test/Dialect/Tensor/canonicalize.mlir

func @fold_extract(%arg0 : index) -> (f32, f16, f16, i32) {
   %ext_4 = tensor.extract %3[%const_1, %const_0, %const_3] : tensor<2x1x4xi32>
 
   // CHECK-NEXT: return [[C4]], [[CM2]], [[C0]], [[C64]]
   return %ext_1, %ext_2, %ext_3, %ext_4 : f32, f16, f16, i32
 }
 
 // -----
 
+// CHECK-LABEL: func @fold_insert
+func @fold_insert(%arg0 : index) -> (tensor<4xf32>) {
+  // Fold an insert into a splat.
+  // CHECK-DAG: %[[C4:.+]] = constant dense<4.{{0*}}e+00> : tensor<4xf32>
+  %0 = constant dense<4.0> : tensor<4xf32>
+  %1 = constant 4.0 : f32
+  %ins_1 = tensor.insert %1 into %0[%arg0] : tensor<4xf32>
+  // CHECK-NEXT: return %[[C4]]
+  return %ins_1 : tensor<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @extract_from_tensor.cast
 // CHECK-SAME: %[[TENSOR:.*]]: tensor<*xf32>
 func @extract_from_tensor.cast(%tensor: tensor<*xf32>) -> f32 {
   // CHECK-NEXT: %[[C0:.*]] = constant 0 : index
   %c0 = constant 0 : index
   // CHECK-NOT: tensor.cast
   %casted = tensor.cast %tensor : tensor<*xf32> to tensor<?xf32>
   // CHECK-NEXT: tensor.extract %[[TENSOR]][%[[C0]]]

mlir/test/Dialect/Tensor/invalid.mlir

 func @extract_too_many_indices(%arg0: tensor<?xf32>) {
   // expected-error@+1 {{incorrect number of indices for extract_element}}
   %0 = tensor.extract %arg0[] : tensor<?xf32>
   return
 }
 
 // -----
 
+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>'}}
   %c0 = constant 0 : i32
   %0 = tensor.from_elements %c0 : tensor<*xi32>
   return
 }
 
 // -----

mlir/test/Dialect/Tensor/ops.mlir

 // CHECK-SAME:                  %[[TENSOR:.*]]: tensor<?x?x?xf32>,
 // CHECK-SAME:                  %[[INDEX:.*]]: index) {
 func @extract(%arg0: tensor<?x?x?xf32>, %arg1: index) {
   // CHECK: tensor.extract %[[TENSOR]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?x?x?xf32>
   %0 = tensor.extract %arg0[%arg1, %arg1, %arg1] : tensor<?x?x?xf32>
   return
 }
 
+// CHECK-LABEL:   func @insert(
+// CHECK-SAME:                  %[[SCALAR:.*]]: f32
+// CHECK-SAME:                  %[[INDEX:.*]]: index
+// CHECK-SAME:                  %[[DEST1:.*]]: tensor<?x?x?xf32>
+// CHECK-SAME:                  %[[DEST2:.*]]: tensor<*xf32>
+func @insert(%arg0: f32, %arg1: index, %arg2: tensor<?x?x?xf32>, %arg3: tensor<*xf32>) {
+  // CHECK: tensor.insert %[[SCALAR]] into %[[DEST1]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?x?x?xf32>
+  %0 = tensor.insert %arg0 into %arg2[%arg1, %arg1, %arg1] : tensor<?x?x?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 = "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>