[Mlir-commits] [mlir] [mlir][tensor] Add a tensor.concat operation (PR #72779)

Matthias Springer llvmlistbot at llvm.org
Tue Nov 21 04:47:20 PST 2023


================
@@ -121,6 +121,66 @@ def Tensor_CastOp : Tensor_Op<"cast", [
   let hasCanonicalizer = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// ConcatOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_ConcatOp : Tensor_Op<"concat",
+    [Pure,
+     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+  let summary = "tensor concatenation operation";
+  let description = [{
+    The "concat" operation constructs a tensor out of a variadic list of input
+    tensors, concatenated along a static dimension. All inputs and the result
+    type must share the same rank.
+
+    `dim` specifies the dimension along which to concatenate. The size of the
+    concatenated dimension in the result must be equal to the sum of the sizes
+    of the inputs along that dimension. All other dimensions in both the inputs
+    and result must be the same size.
+
+    Example:
+
+    ```mlir
+    %0 = tensor.concat dim(0) %0, %1, %2 :
----------------
matthias-springer wrote:

No strong opinion from my side regarding DPS/non-DPS. For the sake of simplicity (and to some degree consistency with the tensor dialect; most ops in there are non-DPS), I'd slightly prefer a non-DPS variant that allocates when bufferizing it. We can lower the op to `tensor.empty` + multiple `tensor.insert_slice`, then use empty tensor elimination to infer the destination. Basically the same way we handle `tensor.pad` today.

I've been thinking of extending `bufferization.materialize_in_destination` with a region that can wrap arbitrary non-DPS op and provide the missing destination. That could be an alternative to turning more and more ops into DPS style.

Regarding `memref.concat`: I would just generate `memref.subview`+`memref.copy` directly to keep the memref dialect as simple as possible. (Unless you want to do tiling of memref IR.)


https://github.com/llvm/llvm-project/pull/72779


More information about the Mlir-commits mailing list