diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h @@ -331,16 +331,19 @@ AffineLoops = 1, ParallelLoops = 2, }; + using TileSizeComputationFunction = std::function(OpBuilder &, Operation *)>; + struct LinalgTilingOptions { /// Computation function that returns the tile sizes for each operation. /// Delayed construction of constant tile sizes should occur to interoperate /// with folding. TileSizeComputationFunction tileSizeComputationFunction = nullptr; + LinalgTilingOptions & - setTileSizeComputationFunction(TileSizeComputationFunction &fun) { - tileSizeComputationFunction = fun; + setTileSizeComputationFunction(TileSizeComputationFunction fun) { + tileSizeComputationFunction = std::move(fun); return *this; } /// Set the `tileSizeComputationFunction` to return the values `ts`. The @@ -356,13 +359,16 @@ LinalgTilingOptions &setTileSizes(ArrayRef ts); /// The interchange vector to reorder the tiled loops. - SmallVector interchangeVector{}; + SmallVector interchangeVector = {}; + LinalgTilingOptions &setInterchange(ArrayRef interchange) { interchangeVector.assign(interchange.begin(), interchange.end()); return *this; } + /// The type of tile loops to generate. - LinalgTilingLoopType loopType{LinalgTilingLoopType::Loops}; + LinalgTilingLoopType loopType = LinalgTilingLoopType::Loops; + LinalgTilingOptions &setLoopType(LinalgTilingLoopType lt) { loopType = lt; return *this; @@ -371,9 +377,10 @@ /// When specified, specifies distribution of generated tile loops to /// processors. Optional distribution = None; + LinalgTilingOptions & - setDistributionOptions(LinalgLoopDistributionOptions &distributionOptions) { - distribution = distributionOptions; + setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) { + distribution = std::move(distributionOptions); return *this; } };