Skip to content

Commit

Permalink
#1772: Added support for remaining ttnn pipeline pybinds and refactor…
Browse files Browse the repository at this point in the history 
…ed naming convention for ttnn pipeline passes
  • Loading branch information
@tapspatel
tapspatel committed Jan 14, 2025
1 parent 731de13 commit 6dd9508
Show file tree
Hide file tree
Showing 3 changed files with 375 additions and 75 deletions.
79 changes: 66 additions & 13 deletions include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,38 +143,91 @@ struct TTIRToTTNNBackendPipelineOptions
//
struct TTIRToEmitCPipelineOptions : public TTIRToTTNNBackendPipelineOptions {};

void createTTNNPipelineTTIRPasses(
void createTTNNOptimizerHelper(OpPassManager &pm,
const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNWorkaroundsHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNDeallocateHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTIRImplicitDeviceHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNPipelineAnalysisPasses(
void createTTIRToTTIRDecompositionPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createInlinerPassHelper(OpPassManager &pm,
const TTIRToTTNNBackendPipelineOptions &options);

void createTTIRLoadSystemDescHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNLayoutHelper(OpPassManager &pm,
const TTIRToTTNNBackendPipelineOptions &options);

void createConvertTTIRToTTNNPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createRemoveDeadValuesPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createCanonicalizerPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNDecomposeLayoutsHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNPipelineLoweringPasses(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNPipelineLayoutDecompositionPass(
void createTTNNPipelineTTIRPasses(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNPipelineDeallocPass(
void createTTNNPipelineWorkaroundPass(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNPipelineTTIRPassesFromString(OpPassManager &pm,
std::string options);
void createTTIRToTTNNBackendPipeline(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);

void createTTNNPipelineAnalysisPassesFromString(OpPassManager &pm,
void createTTNNDecomposeLayoutsHelperFromString(OpPassManager &pm,
std::string options);

void createTTNNDeallocateHelperFromString(OpPassManager &pm,
std::string options);

void createTTIRToTTIRDecompositionPassFromString(OpPassManager &pm,
std::string options);

void createInlinerPassFromString(OpPassManager &pm, std::string options);

void createTTIRLoadSystemDescFromString(OpPassManager &pm, std::string options);

void createTTIRImplicitDeviceFromString(OpPassManager &pm, std::string options);

void createTTNNLayoutFromString(OpPassManager &pm, std::string options);

void createConvertTTIRToTTNNPassFromString(OpPassManager &pm,
std::string options);

void createRemoveDeadValuesPassFromString(OpPassManager &pm,
std::string options);

void createTTNNPipelineLoweringPassesFromString(OpPassManager &pm,
std::string options);

void createTTNNPipelineLayoutDecompositionPassFromString(OpPassManager &pm,
std::string options);
void createTTNNPipelineTTIRPassesFromString(OpPassManager &pm,
std::string options);

void createTTNNPipelineDeallocPassFromString(OpPassManager &pm,
std::string options);
void createTTNNWorkaroundsFromString(OpPassManager &pm, std::string options);

void createTTIRToTTNNBackendPipeline(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options);
void createCanonicalizerPassFromString(OpPassManager &pm, std::string options);

void createTTNNPipelineWorkaroundPassFromString(OpPassManager &pm,
std::string options);

void createTTNNOptimizerFromString(OpPassManager &pm, std::string options);

void createTTIRToEmitCPipeline(OpPassManager &pm,
const TTIRToEmitCPipelineOptions &options);
Expand Down
214 changes: 161 additions & 53 deletions lib/Dialect/TTNN/Pipelines/TTNNPipelines.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,26 +17,7 @@ namespace mlir::tt::ttnn {
// Pipeline implementation.
//===----------------------------------------------------------------------===//

void createTTNNPipelineTTIRPasses(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
ttir::TTIRLoadSystemDescOptions systemDescOptions;
systemDescOptions.path = options.systemDescPath;

pm.addPass(mlir::tt::createTTIRToTTIRDecompositionPass());

// Inlines all private functions. I.e flattens the program into the main
// function. Removes all private functions.
pm.addPass(mlir::createInlinerPass());

pm.addPass(mlir::tt::ttir::createTTIRLoadSystemDesc(systemDescOptions));

ttir::TTIRImplicitDeviceOptions implicitDeviceOptions;
implicitDeviceOptions.meshShape = ::llvm::SmallVector<int64_t>(
options.meshShape.begin(), options.meshShape.end());
pm.addPass(mlir::tt::ttir::createTTIRImplicitDevice(implicitDeviceOptions));
}

void createTTNNPipelineAnalysisPasses(
void createTTNNOptimizerHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
if (options.optimizerPassEnabled) {
ttnn::TTNNOptimizerOptions optimizerOptions;
Expand All @@ -53,48 +34,167 @@ void createTTNNPipelineAnalysisPasses(
}
}

void createTTNNPipelineLoweringPasses(
void createTTNNWorkaroundsHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
TTNNWorkaroundsOptions workaroundOptions{
options.layouotWorkaroundsEnabled,
options.decompositionWorkaroundsEnabled};
pm.addPass(createTTNNWorkarounds(workaroundOptions));
}

void createTTNNDeallocateHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
// Add pass to add layout information.
pm.addPass(createTTNNDeallocate());
}

void createTTIRImplicitDeviceHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
ttir::TTIRImplicitDeviceOptions implicitDeviceOptions;
implicitDeviceOptions.meshShape = ::llvm::SmallVector<int64_t>(
options.meshShape.begin(), options.meshShape.end());
pm.addPass(mlir::tt::ttir::createTTIRImplicitDevice(implicitDeviceOptions));
}

void createTTIRToTTIRDecompositionPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
pm.addPass(mlir::tt::createTTIRToTTIRDecompositionPass());
}

void createInlinerPassHelper(OpPassManager &pm,
const TTIRToTTNNBackendPipelineOptions &options) {
// Inlines all private functions. I.e flattens the program into the main
// function. Removes all private functions.
pm.addPass(mlir::createInlinerPass());
}

void createTTIRLoadSystemDescHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
ttir::TTIRLoadSystemDescOptions systemDescOptions;
systemDescOptions.path = options.systemDescPath;
pm.addPass(mlir::tt::ttir::createTTIRLoadSystemDesc(systemDescOptions));
}

void createTTNNLayoutHelper(OpPassManager &pm,
const TTIRToTTNNBackendPipelineOptions &options) {
pm.addPass(createTTNNLayout());
// Add pass to convert TTIR to TTNN.
}

void createConvertTTIRToTTNNPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
pm.addPass(createConvertTTIRToTTNNPass());
// Add pass to remove unused values.
}

void createRemoveDeadValuesPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
pm.addPass(mlir::createRemoveDeadValuesPass());
}

// Create a pass to workaround issues in the TTNN dialect.
void createTTNNPipelineWorkaroundPass(
void createCanonicalizerPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
TTNNWorkaroundsOptions workaroundOptions{
options.layouotWorkaroundsEnabled,
options.decompositionWorkaroundsEnabled};
pm.addPass(createTTNNWorkarounds(workaroundOptions));
pm.addPass(mlir::createCanonicalizerPass());
}

void createTTNNPipelineLayoutDecompositionPass(
void createConvertTTNNToEmitCPassHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
pm.addPass(createConvertTTNNToEmitCPass());
}

void createTTNNDecomposeLayoutsHelper(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
pm.addPass(createTTNNDecomposeLayouts());
}

void createTTNNPipelineDeallocPass(
void createTTNNPipelineLoweringPasses(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
createTTNNLayoutHelper(pm, options);
createConvertTTIRToTTNNPassHelper(pm, options);
createRemoveDeadValuesPassHelper(pm, options);
}

void createTTNNPipelineTTIRPasses(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
pm.addPass(createTTNNDeallocate());
createTTIRToTTIRDecompositionPassHelper(pm, options);
createInlinerPassHelper(pm, options);
createTTIRLoadSystemDescHelper(pm, options);
createTTIRImplicitDeviceHelper(pm, options);
}

void createTTNNPipelineTTIRPassesFromString(OpPassManager &pm,
std::string options) {
// Create a pass to workaround issues in the TTNN dialect.
void createTTNNPipelineWorkaroundPass(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
createTTNNWorkaroundsHelper(pm, options);
createCanonicalizerPassHelper(pm, options);
}

void createTTIRToTTNNBackendPipeline(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
createTTNNPipelineTTIRPasses(pm, options);
createTTNNPipelineLoweringPasses(pm, options);
createTTNNPipelineWorkaroundPass(pm, options);
createTTNNOptimizerHelper(pm, options);
createTTNNDecomposeLayoutsHelper(pm, options);
createTTNNDeallocateHelper(pm, options);
}

void createTTNNDecomposeLayoutsHelperFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNPipelineTTIRPasses(pm, *optionsStruct);
createTTNNDecomposeLayoutsHelper(pm, *optionsStruct);
}

void createTTNNPipelineAnalysisPassesFromString(OpPassManager &pm,
std::string options) {
void createTTNNDeallocateHelperFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNPipelineAnalysisPasses(pm, *optionsStruct);
createTTNNDeallocateHelper(pm, *optionsStruct);
}

void createTTIRToTTIRDecompositionPassFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTIRToTTIRDecompositionPassHelper(pm, *optionsStruct);
}

void createInlinerPassFromString(OpPassManager &pm, std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createInlinerPassHelper(pm, *optionsStruct);
}

void createTTIRLoadSystemDescFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTIRLoadSystemDescHelper(pm, *optionsStruct);
}

void createTTIRImplicitDeviceFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTIRImplicitDeviceHelper(pm, *optionsStruct);
}

void createTTNNLayoutFromString(OpPassManager &pm, std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNLayoutHelper(pm, *optionsStruct);
}

void createConvertTTIRToTTNNPassFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createConvertTTIRToTTNNPassHelper(pm, *optionsStruct);
}

void createRemoveDeadValuesPassFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createRemoveDeadValuesPassHelper(pm, *optionsStruct);
}

void createTTNNPipelineLoweringPassesFromString(OpPassManager &pm,
Expand All @@ -104,28 +204,36 @@ void createTTNNPipelineLoweringPassesFromString(OpPassManager &pm,
createTTNNPipelineLoweringPasses(pm, *optionsStruct);
}

void createTTNNPipelineLayoutDecompositionPassFromString(OpPassManager &pm,
std::string options) {
void createTTNNPipelineTTIRPassesFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNPipelineTTIRPasses(pm, *optionsStruct);
}

void createTTNNWorkaroundsFromString(OpPassManager &pm, std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNPipelineLayoutDecompositionPass(pm, *optionsStruct);
createTTNNWorkaroundsHelper(pm, *optionsStruct);
}

void createTTNNPipelineDeallocPassFromString(OpPassManager &pm,
std::string options) {
void createCanonicalizerPassFromString(OpPassManager &pm, std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNPipelineDeallocPass(pm, *optionsStruct);
createCanonicalizerPassHelper(pm, *optionsStruct);
}

void createTTIRToTTNNBackendPipeline(
OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
createTTNNPipelineTTIRPasses(pm, options);
createTTNNPipelineLoweringPasses(pm, options);
createTTNNPipelineWorkaroundPass(pm, options);
createTTNNPipelineAnalysisPasses(pm, options);
createTTNNPipelineLayoutDecompositionPass(pm, options);
createTTNNPipelineDeallocPass(pm, options);
void createTTNNPipelineWorkaroundPassFromString(OpPassManager &pm,
std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNPipelineWorkaroundPass(pm, *optionsStruct);
}

void createTTNNOptimizerFromString(OpPassManager &pm, std::string options) {
auto optionsStruct =
TTIRToTTNNBackendPipelineOptions::createFromString(options);
createTTNNOptimizerHelper(pm, *optionsStruct);
}

void createTTIRToEmitCPipeline(OpPassManager &pm,
Expand Down
Loading

0 comments on commit 6dd9508

Please sign in to comment.