diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td --- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td @@ -1827,6 +1827,13 @@ `into` $dest attr-dict `:` type($source) `->` type($dest) }]; + let builders = [ + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$innerDimsPos, + "ArrayRef":$innerTiles, + CArg<"ArrayRef", "{}">:$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 --- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -3450,6 +3450,24 @@ return Speculation::Speculatable; } +void UnPackOp::build(OpBuilder &builder, OperationState &state, Value source, + Value dest, ArrayRef innerDimsPos, + ArrayRef innerTiles, + ArrayRef outerDimsPerm) { + assert(innerDimsPos.size() == innerTiles.size() && + "number of tile sizes specified must match the specified number of " + "original dimensions to be tiled"); + SmallVector staticTileSizes; + SmallVector 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) {