54 lines · cpp
1//===- TestSlice.cpp - Test slice related analisis ------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Analysis/TopologicalSortUtils.h"10#include "mlir/IR/BuiltinTypes.h"11#include "mlir/IR/SymbolTable.h"12#include "mlir/Pass/Pass.h"13 14using namespace mlir;15 16static const StringLiteral kToSortMark = "test_to_sort";17static const StringLiteral kOrderIndex = "test_sort_index";18 19namespace {20 21struct TestTopologicalSortPass22 : public PassWrapper<TestTopologicalSortPass,23 InterfacePass<SymbolOpInterface>> {24 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTopologicalSortPass)25 26 StringRef getArgument() const final { return "test-print-topological-sort"; }27 StringRef getDescription() const final {28 return "Sorts operations topologically and attaches attributes with their "29 "corresponding index in the ordering to them";30 }31 void runOnOperation() override {32 SetVector<Operation *> toSort;33 getOperation().walk([&](Operation *op) {34 if (op->hasAttrOfType<UnitAttr>(kToSortMark))35 toSort.insert(op);36 });37 38 auto i32Type = IntegerType::get(&getContext(), 32);39 SetVector<Operation *> sortedOps = topologicalSort(toSort);40 for (auto [index, op] : llvm::enumerate(sortedOps))41 op->setAttr(kOrderIndex, IntegerAttr::get(i32Type, index));42 }43};44 45} // namespace46 47namespace mlir {48namespace test {49void registerTestSliceAnalysisPass() {50 PassRegistration<TestTopologicalSortPass>();51}52} // namespace test53} // namespace mlir54