diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h index 493180cd54e5b..a8cc37a103f04 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h @@ -2,6 +2,7 @@ #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H #include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h" +#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td index 3bbb8b02c644e..f53f569070f09 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td @@ -284,6 +284,15 @@ def EmptyTensorToAllocTensorPass : Pass<"empty-tensor-to-alloc-tensor"> { let dependentDialects = ["tensor::TensorDialect"]; } +def layoutMapClValues { + string values = [{ + ::llvm::cl::values( + clEnumValN(LayoutMapOption::InferLayoutMap, "infer-layout-map", ""), + clEnumValN(LayoutMapOption::IdentityLayoutMap, "identity-layout-map", ""), + clEnumValN(LayoutMapOption::FullyDynamicLayoutMap, "fully-dynamic-layout-map", "") + )}]; +} + def OneShotBufferizePass : Pass<"one-shot-bufferize", "ModuleOp"> { let summary = "One-Shot Bufferize"; let description = [{ @@ -424,9 +433,10 @@ def OneShotBufferizePass : Pass<"one-shot-bufferize", "ModuleOp"> { "Skip analysis of functions with these symbol names." "Set copyBeforeWrite to true when bufferizing them.">, Option<"functionBoundaryTypeConversion", - "function-boundary-type-conversion", "std::string", - /*default=*/"\"infer-layout-map\"", - "Controls layout maps when bufferizing function signatures.">, + "function-boundary-type-conversion", "LayoutMapOption", + /*default=*/"LayoutMapOption::InferLayoutMap", + "Controls layout maps when bufferizing function signatures.", + layoutMapClValues.values>, Option<"mustInferMemorySpace", "must-infer-memory-space", "bool", /*default=*/"false", "The memory space of an memref types must always be inferred. If " @@ -444,9 +454,10 @@ def OneShotBufferizePass : Pass<"one-shot-bufferize", "ModuleOp"> { /*default=*/"false", "Test only: Annotate IR with RaW conflicts. Requires " "test-analysis-only.">, - Option<"unknownTypeConversion", "unknown-type-conversion", "std::string", - /*default=*/"\"fully-dynamic-layout-map\"", - "Controls layout maps for non-inferrable memref types.">, + Option<"unknownTypeConversion", "unknown-type-conversion", "LayoutMapOption", + /*default=*/"LayoutMapOption::FullyDynamicLayoutMap", + "Controls layout maps for non-inferrable memref types.", + layoutMapClValues.values>, Option<"bufferAlignment", "buffer-alignment", "uint64_t", /*default=*/"64", "Sets the alignment of newly allocated buffers.">, diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp index c0e3fca428376..e97b34b20ff72 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -38,16 +38,6 @@ using namespace mlir::bufferization; namespace { -static LayoutMapOption parseLayoutMapOption(const std::string &s) { - if (s == "fully-dynamic-layout-map") - return LayoutMapOption::FullyDynamicLayoutMap; - if (s == "identity-layout-map") - return LayoutMapOption::IdentityLayoutMap; - if (s == "infer-layout-map") - return LayoutMapOption::InferLayoutMap; - llvm_unreachable("invalid layout map option"); -} - static OneShotBufferizationOptions::AnalysisHeuristic parseHeuristicOption(const std::string &s) { if (s == "bottom-up") @@ -83,8 +73,7 @@ struct OneShotBufferizePass opt.analysisHeuristic = parseHeuristicOption(analysisHeuristic); opt.copyBeforeWrite = copyBeforeWrite; opt.dumpAliasSets = dumpAliasSets; - opt.setFunctionBoundaryTypeConversion( - parseLayoutMapOption(functionBoundaryTypeConversion)); + opt.setFunctionBoundaryTypeConversion(functionBoundaryTypeConversion); if (mustInferMemorySpace && useEncodingForMemorySpace) { emitError(getOperation()->getLoc()) @@ -118,8 +107,7 @@ struct OneShotBufferizePass opt.noAnalysisFuncFilter = noAnalysisFuncFilter; // Configure type converter. - LayoutMapOption unknownTypeConversionOption = - parseLayoutMapOption(unknownTypeConversion); + LayoutMapOption unknownTypeConversionOption = unknownTypeConversion; if (unknownTypeConversionOption == LayoutMapOption::InferLayoutMap) { emitError(UnknownLoc::get(&getContext()), "Invalid option: 'infer-layout-map' is not a valid value for "