diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/Patterns.h b/mlir/include/mlir/Dialect/SCF/Transforms/Patterns.h --- a/mlir/include/mlir/Dialect/SCF/Transforms/Patterns.h +++ b/mlir/include/mlir/Dialect/SCF/Transforms/Patterns.h @@ -14,31 +14,11 @@ #include "mlir/IR/PatternMatch.h" namespace mlir { + +class ConversionTarget; +class TypeConverter; + namespace scf { -/// Generate a pipelined version of the scf.for loop based on the schedule given -/// as option. This applies the mechanical transformation of changing the loop -/// and generating the prologue/epilogue for the pipelining and doesn't make any -/// decision regarding the schedule. -/// Based on the options the loop is split into several stages. -/// The transformation assumes that the scheduling given by user is valid. -/// For example if we break a loop into 3 stages named S0, S1, S2 we would -/// generate the following code with the number in parenthesis as the iteration -/// index: -/// -/// S0(0) // Prologue -/// S0(1) S1(0) // Prologue -/// scf.for %I = %C0 to %N - 2 { -/// S0(I+2) S1(I+1) S2(I) // Pipelined kernel -/// } -/// S1(N) S2(N-1) // Epilogue -/// S2(N) // Epilogue -/// -/// If `modifiedIR` is provided, it will be set to a value that indicates -/// whether pipelining modified the IR before failing, signaling to the caller -/// whether they can proceed with different transformations. -FailureOr pipelineForLoop(RewriterBase &rewriter, ForOp forOp, - const PipeliningOption &options, - bool *modifiedIR = nullptr); // TODO: such patterns should be auto-generated. class ForLoopPipeliningPattern : public OpRewritePattern { @@ -60,6 +40,35 @@ PipeliningOption options; }; +/// Populates patterns for SCF structural type conversions and sets up the +/// provided ConversionTarget with the appropriate legality configuration for +/// the ops to get converted properly. +/// +/// A "structural" type conversion is one where the underlying ops are +/// completely agnostic to the actual types involved and simply need to update +/// their types. An example of this is scf.if -- the scf.if op and the +/// corresponding scf.yield ops need to update their types accordingly to the +/// TypeConverter, but otherwise don't care what type conversions are happening. +void populateSCFStructuralTypeConversionsAndLegality( + TypeConverter &typeConverter, RewritePatternSet &patterns, + ConversionTarget &target); + +/// Populates the provided pattern set with patterns that do 1:N type +/// conversions on (some) SCF ops. This is intended to be used with +/// applyPartialOneToNConversion. +void populateSCFStructuralOneToNTypeConversions(TypeConverter &typeConverter, + RewritePatternSet &patterns); + +/// Populate patterns for SCF software pipelining transformation. See the +/// ForLoopPipeliningPattern for the transformation details. +void populateSCFLoopPipeliningPatterns(RewritePatternSet &patterns, + const PipeliningOption &options); + +/// Populate patterns for canonicalizing operations inside SCF loop bodies. +/// At the moment, only affine.min/max computations with iteration variables, +/// loop bounds and loop steps are canonicalized. +void populateSCFForLoopCanonicalizationPatterns(RewritePatternSet &patterns); + } // namespace scf } // namespace mlir diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h @@ -19,18 +19,11 @@ namespace mlir { -class AffineMap; -class ConversionTarget; struct LogicalResult; -class MLIRContext; class Region; class RewriterBase; -class TypeConverter; -class RewritePatternSet; class Operation; class Value; -class ValueRange; -class PatternRewriter; namespace scf { @@ -107,25 +100,6 @@ tileParallelLoop(ParallelOp op, llvm::ArrayRef tileSizes, bool noMinMaxBounds); -/// Populates patterns for SCF structural type conversions and sets up the -/// provided ConversionTarget with the appropriate legality configuration for -/// the ops to get converted properly. -/// -/// A "structural" type conversion is one where the underlying ops are -/// completely agnostic to the actual types involved and simply need to update -/// their types. An example of this is scf.if -- the scf.if op and the -/// corresponding scf.yield ops need to update their types accordingly to the -/// TypeConverter, but otherwise don't care what type conversions are happening. -void populateSCFStructuralTypeConversionsAndLegality( - TypeConverter &typeConverter, RewritePatternSet &patterns, - ConversionTarget &target); - -/// Populates the provided pattern set with patterns that do 1:N type -/// conversions on (some) SCF ops. This is intended to be used with -/// applyPartialOneToNConversion. -void populateSCFStructuralOneToNTypeConversions(TypeConverter &typeConverter, - RewritePatternSet &patterns); - /// Options to dictate how loops should be pipelined. struct PipeliningOption { /// Lambda returning all the operation in the forOp, with their stage, in the @@ -167,15 +141,30 @@ // TODO: add option to decide if the prologue should be peeled. }; -/// Populate patterns for SCF software pipelining transformation. See the -/// ForLoopPipeliningPattern for the transformation details. -void populateSCFLoopPipeliningPatterns(RewritePatternSet &patterns, - const PipeliningOption &options); - -/// Populate patterns for canonicalizing operations inside SCF loop bodies. -/// At the moment, only affine.min/max computations with iteration variables, -/// loop bounds and loop steps are canonicalized. -void populateSCFForLoopCanonicalizationPatterns(RewritePatternSet &patterns); +/// Generate a pipelined version of the scf.for loop based on the schedule given +/// as option. This applies the mechanical transformation of changing the loop +/// and generating the prologue/epilogue for the pipelining and doesn't make any +/// decision regarding the schedule. +/// Based on the options the loop is split into several stages. +/// The transformation assumes that the scheduling given by user is valid. +/// For example if we break a loop into 3 stages named S0, S1, S2 we would +/// generate the following code with the number in parenthesis as the iteration +/// index: +/// +/// S0(0) // Prologue +/// S0(1) S1(0) // Prologue +/// scf.for %I = %C0 to %N - 2 { +/// S0(I+2) S1(I+1) S2(I) // Pipelined kernel +/// } +/// S1(N) S2(N-1) // Epilogue +/// S2(N) // Epilogue +/// +/// If `modifiedIR` is provided, it will be set to a value that indicates +/// whether pipelining modified the IR before failing, signaling to the caller +/// whether they can proceed with different transformations. +FailureOr pipelineForLoop(RewriterBase &rewriter, ForOp forOp, + const PipeliningOption &options, + bool *modifiedIR = nullptr); } // namespace scf } // namespace mlir diff --git a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp --- a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp +++ b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp @@ -17,7 +17,6 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h" #include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Dialect/SCF/Transforms/Transforms.h" #include "mlir/Dialect/Utils/IndexingUtils.h" #include "mlir/Dialect/Vector/IR/VectorOps.h" diff --git a/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp b/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp @@ -12,7 +12,7 @@ #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/SCF/Transforms/Transforms.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Transforms/DialectConversion.h" namespace mlir { diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp --- a/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp @@ -16,7 +16,7 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/SCF/Transforms/Transforms.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/PatternMatch.h" diff --git a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp --- a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp @@ -7,8 +7,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/SCF/Transforms/Passes.h" -#include "mlir/Dialect/SCF/Transforms/Transforms.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Transforms/DialectConversion.h" #include diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp @@ -15,7 +15,7 @@ #include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/SCF/Transforms/Transforms.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" #include "mlir/Dialect/SparseTensor/Transforms/Passes.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" diff --git a/mlir/test/lib/Conversion/OneToNTypeConversion/TestOneToNTypeConversionPass.cpp b/mlir/test/lib/Conversion/OneToNTypeConversion/TestOneToNTypeConversionPass.cpp --- a/mlir/test/lib/Conversion/OneToNTypeConversion/TestOneToNTypeConversionPass.cpp +++ b/mlir/test/lib/Conversion/OneToNTypeConversion/TestOneToNTypeConversionPass.cpp @@ -8,7 +8,7 @@ #include "TestDialect.h" #include "mlir/Dialect/Func/Transforms/OneToNFuncConversions.h" -#include "mlir/Dialect/SCF/Transforms/Transforms.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/OneToNTypeConversion.h" diff --git a/mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp b/mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp --- a/mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp +++ b/mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp @@ -13,7 +13,7 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/SCF/Transforms/Transforms.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Pass/Pass.h" #include "mlir/Pass/PassManager.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" diff --git a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp --- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp +++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp @@ -14,7 +14,7 @@ #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/SCF/Transforms/Transforms.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/Dialect/SCF/Utils/Utils.h" #include "mlir/IR/Builders.h" #include "mlir/IR/PatternMatch.h"