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 @@ -1761,6 +1761,14 @@ `into` $dest attr-dict `:` type($source) `->` type($dest) }]; + let builders = [ + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$innerDimsPos, + "ArrayRef":$innerTiles, + CArg<"Optional", "llvm::None">:$paddingValue, + CArg<"ArrayRef", "{}">:$outerDimsPerm)> + ]; + let extraClassDeclaration = commonExtraClassDeclaration # [{ // Method to get the `ShapedType` of the result based on the inner tiles, // position of the inner tiles (innerDimsPos) and interchange vector of 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 @@ -3238,6 +3238,26 @@ setNameFn(getResult(), "pack"); } +void PackOp::build(OpBuilder &builder, OperationState &state, Value source, + Value dest, ArrayRef innerDimsPos, + ArrayRef innerTiles, + Optional paddingValue, + 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, + paddingValue ? paddingValue.value() : nullptr, + outerDimsPerm.empty() ? nullptr + : builder.getDenseI64ArrayAttr(outerDimsPerm), + builder.getDenseI64ArrayAttr(innerDimsPos), dynamicTileSizes, + builder.getDenseI64ArrayAttr(staticTileSizes)); +} + LogicalResult PackOp::reifyResultShapes(OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {