[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