mlir/test/lib/Analysis/TestMatchReduction.cpp


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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

//===- TestMatchReduction.cpp - Test the match reduction utility ----------===//
//
// 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 contains a test pass for the match reduction utility.
//
//===----------------------------------------------------------------------===//

#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/Pass.h"

using namespace mlir;

namespace {

void printReductionResult(Operation *redRegionOp, unsigned numOutput,
                          Value reducedValue,
                          ArrayRef<Operation *> combinerOps) {
  if (reducedValue) {
    redRegionOp->emitRemark("Reduction found in output #") << numOutput << "!";
    redRegionOp->emitRemark("Reduced Value: ") << reducedValue;
    for (Operation *combOp : combinerOps)
      redRegionOp->emitRemark("Combiner Op: ") << *combOp;

    return;
  }

  redRegionOp->emitRemark("Reduction NOT found in output #")
      << numOutput << "!";
}

struct TestMatchReductionPass
    : public PassWrapper<TestMatchReductionPass,
                         InterfacePass<FunctionOpInterface>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMatchReductionPass)

  StringRef getArgument() const final { return "test-match-reduction"; }
  StringRef getDescription() const final {
    return "Test the match reduction utility.";
  }

  void runOnOperation() override {
    FunctionOpInterface func = getOperation();
    func->emitRemark("Testing function");

    func.walk<WalkOrder::PreOrder>([](Operation *op) {
      if (isa<FunctionOpInterface>(op))
        return;

      // Limit testing to ops with only one region.
      if (op->getNumRegions() != 1)
        return;

      Region &region = op->getRegion(0);
      if (!region.hasOneBlock())
        return;

      // We expect all the tested region ops to have 1 input by default. The
      // remaining arguments are assumed to be outputs/reductions and there must
      // be at least one.
      // TODO: Extend it to support more generic cases.
      Block &regionEntry = region.front();
      auto args = regionEntry.getArguments();
      if (args.size() < 2)
        return;

      auto outputs = args.drop_front();
      for (int i = 0, size = outputs.size(); i < size; ++i) {
        SmallVector<Operation *, 4> combinerOps;
        Value reducedValue = matchReduction(outputs, i, combinerOps);
        printReductionResult(op, i, reducedValue, combinerOps);
      }
    });
  }
};

} // namespace

namespace mlir {
namespace test {
void registerTestMatchReductionPass() {
  PassRegistration<TestMatchReductionPass>();
}
} // namespace test
} // namespace mlir