mlir/test/lib/Dialect/Transform/TestTransformStateExtension.h


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

//===- TestTransformStateExtension.h - Test Utility -------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file defines an TransformState extension for the purpose of testing the
// relevant APIs.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_TEST_LIB_DIALECT_TRANSFORM_TESTTRANSFORMSTATEEXTENSION_H
#define MLIR_TEST_LIB_DIALECT_TRANSFORM_TESTTRANSFORMSTATEEXTENSION_H

#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"

using namespace mlir;

namespace mlir {
namespace test {
class TestTransformStateExtension
    : public transform::TransformState::Extension {
public:
  TestTransformStateExtension(transform::TransformState &state,
                              StringAttr message)
      : Extension(state), message(message) {}

  StringRef getMessage() const { return message.getValue(); }

  LogicalResult updateMapping(Operation *previous, Operation *updated);

private:
  StringAttr message;
};

class TransformStateInitializerExtension
    : public transform::TransformState::Extension {
public:
  TransformStateInitializerExtension(transform::TransformState &state,
                                     int numOp,
                                     SmallVector<std::string> &registeredOps)
      : Extension(state), numOp(numOp), registeredOps(registeredOps) {}

  int getNumOp() { return numOp; }
  void setNumOp(int num) { numOp = num; }
  SmallVector<std::string> getRegisteredOps() { return registeredOps; }
  void pushRegisteredOps(const std::string &newOp) {
    registeredOps.push_back(newOp);
  }
  std::string printMessage() const {
    std::string message = "Registered transformOps are: ";
    for (const auto &op : registeredOps) {
      message += op + " | ";
    }
    return message;
  }

private:
  int numOp;
  SmallVector<std::string> registeredOps;
};

} // namespace test
} // namespace mlir

#endif // MLIR_TEST_LIB_DIALECT_TRANSFORM_TESTTRANSFORMSTATEEXTENSION_H