[Mlir-commits] [mlir] 87ecf9d - [MLIR][Tensor] Add custom builder for unpack op

[Lorenzo Chelini]

Author: Lorenzo Chelini
Date: 2022-12-07T12:40:45+01:00
New Revision: 87ecf9d155936651dfff111855bebc5a83a7a3cb

URL: https://github.com/llvm/llvm-project/commit/87ecf9d155936651dfff111855bebc5a83a7a3cb
DIFF: https://github.com/llvm/llvm-project/commit/87ecf9d155936651dfff111855bebc5a83a7a3cb.diff

LOG: [MLIR][Tensor] Add custom builder for unpack op

Reviewed By: hanchung

Differential Revision: https://reviews.llvm.org/D139344

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
    mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index f41bd7d7c563f..55f3695c4d1b8 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1827,6 +1827,13 @@ def Tensor_UnPackOp : Tensor_RelayoutOp<"unpack"> {
     `into` $dest attr-dict `:` type($source) `->` type($dest)
   }];
 
+  let builders = [
+    OpBuilder<(ins "Value":$source, "Value":$dest,
+    "ArrayRef<int64_t>":$innerDimsPos,
+    "ArrayRef<OpFoldResult>":$innerTiles,
+    CArg<"ArrayRef<int64_t>", "{}">:$outerDimsPerm)>
+  ];
+
   let extraClassDeclaration = commonExtraClassDeclaration;
   let hasCanonicalizeMethod = 1;
 }

diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 8faf6cc2e1c8d..e0f45b7b99edc 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3450,6 +3450,24 @@ Speculation::Speculatability UnPackOp::getSpeculatability() {
   return Speculation::Speculatable;
 }
 
+void UnPackOp::build(OpBuilder &builder, OperationState &state, Value source,
+                     Value dest, ArrayRef<int64_t> innerDimsPos,
+                     ArrayRef<OpFoldResult> innerTiles,
+                     ArrayRef<int64_t> outerDimsPerm) {
+  assert(innerDimsPos.size() == innerTiles.size() &&
+         "number of tile sizes specified must match the specified number of "
+         "original dimensions to be tiled");
+  SmallVector<int64_t> staticTileSizes;
+  SmallVector<Value> dynamicTileSizes;
+  dispatchIndexOpFoldResults(innerTiles, dynamicTileSizes, staticTileSizes,
+                             ShapedType::kDynamic);
+  build(builder, state, dest.getType(), source, dest,
+        outerDimsPerm.empty() ? nullptr
+                              : builder.getDenseI64ArrayAttr(outerDimsPerm),
+        builder.getDenseI64ArrayAttr(innerDimsPos), dynamicTileSizes,
+        builder.getDenseI64ArrayAttr(staticTileSizes));
+}
+
 /// pack(unpack(x)) -> x
 LogicalResult UnPackOp::canonicalize(UnPackOp unpackOp,
                                      PatternRewriter &rewriter) {