mlir/test/lib/Analysis/TestMatchReduction.cpp

   1 //===- TestMatchReduction.cpp - Test the match reduction utility ----------===//
   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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file contains a test pass for the match reduction utility.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "mlir/Analysis/SliceAnalysis.h"
  14 #include "mlir/IR/FunctionInterfaces.h"
  15 #include "mlir/Pass/Pass.h"
  16
  17 using namespace mlir;
  18
  19 namespace {
  20
  21 void printReductionResult(Operation *redRegionOp, unsigned numOutput,
  22                           Value reducedValue,
  23                           ArrayRef<Operation *> combinerOps) {
  24   if (reducedValue) {
  25     redRegionOp->emitRemark("Reduction found in output #") << numOutput << "!";
  26     redRegionOp->emitRemark("Reduced Value: ") << reducedValue;
  27     for (Operation *combOp : combinerOps)
  28       redRegionOp->emitRemark("Combiner Op: ") << *combOp;
  29
  30     return;
  31   }
  32
  33   redRegionOp->emitRemark("Reduction NOT found in output #")
  34       << numOutput << "!";
  35 }
  36
  37 struct TestMatchReductionPass
  38     : public PassWrapper<TestMatchReductionPass,
  39                          InterfacePass<FunctionOpInterface>> {
  40   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMatchReductionPass)
  41
  42   StringRef getArgument() const final { return "test-match-reduction"; }
  43   StringRef getDescription() const final {
  44     return "Test the match reduction utility.";
  45   }
  46
  47   void runOnOperation() override {
  48     FunctionOpInterface func = getOperation();
  49     func->emitRemark("Testing function");
  50
  51     func.walk<WalkOrder::PreOrder>([](Operation *op) {
  52       if (isa<FunctionOpInterface>(op))
  53         return;
  54
  55       // Limit testing to ops with only one region.
  56       if (op->getNumRegions() != 1)
  57         return;
  58
  59       Region &region = op->getRegion(0);
  60       if (!region.hasOneBlock())
  61         return;
  62
  63       // We expect all the tested region ops to have 1 input by default. The
  64       // remaining arguments are assumed to be outputs/reductions and there must
  65       // be at least one.
  66       // TODO: Extend it to support more generic cases.
  67       Block &regionEntry = region.front();
  68       auto args = regionEntry.getArguments();
  69       if (args.size() < 2)
  70         return;
  71
  72       auto outputs = args.drop_front();
  73       for (int i = 0, size = outputs.size(); i < size; ++i) {
  74         SmallVector<Operation *, 4> combinerOps;
  75         Value reducedValue = matchReduction(outputs, i, combinerOps);
  76         printReductionResult(op, i, reducedValue, combinerOps);
  77       }
  78     });
  79   }
  80 };
  81
  82 } // namespace
  83
  84 namespace mlir {
  85 namespace test {
  86 void registerTestMatchReductionPass() {
  87   PassRegistration<TestMatchReductionPass>();
  88 }
  89 } // namespace test
  90 } // namespace mlir