mlir/test/lib/Dialect/Affine/TestAccessAnalysis.cpp

   1 //===- TestAccessAnalysis.cpp - Test affine access analysis 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 implements a pass to test affine access analysis utilities.
  10 //
  11 //===----------------------------------------------------------------------===//
  12 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
  13 #include "mlir/Dialect/Affine/Analysis/Utils.h"
  14 #include "mlir/Dialect/Affine/LoopFusionUtils.h"
  15 #include "mlir/Dialect/Func/IR/FuncOps.h"
  16 #include "mlir/Pass/Pass.h"
  17
  18 #define PASS_NAME "test-affine-access-analysis"
  19
  20 using namespace mlir;
  21 using namespace mlir::affine;
  22
  23 namespace {
  24
  25 struct TestAccessAnalysis
  26     : public PassWrapper<TestAccessAnalysis, OperationPass<func::FuncOp>> {
  27   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAccessAnalysis)
  28
  29   StringRef getArgument() const final { return PASS_NAME; }
  30   StringRef getDescription() const final {
  31     return "Tests affine memory access analysis utility";
  32   }
  33
  34   void runOnOperation() override;
  35 };
  36
  37 } // namespace
  38
  39 /// Gathers all affine load/store ops in loop nest rooted at 'forOp' into
  40 /// 'loadAndStoreOps'.
  41 static void
  42 gatherLoadsAndStores(AffineForOp forOp,
  43                      SmallVectorImpl<Operation *> &loadAndStoreOps) {
  44   forOp.walk([&](Operation *op) {
  45     if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
  46       loadAndStoreOps.push_back(op);
  47   });
  48 }
  49
  50 void TestAccessAnalysis::runOnOperation() {
  51   SmallVector<Operation *> loadStores;
  52   SmallVector<AffineForOp> enclosingOps;
  53   // Go over all top-level affine.for ops and test each contained affine
  54   // access's contiguity along every surrounding loop IV.
  55   for (auto forOp : getOperation().getOps<AffineForOp>()) {
  56     loadStores.clear();
  57     gatherLoadsAndStores(forOp, loadStores);
  58     for (Operation *memOp : loadStores) {
  59       enclosingOps.clear();
  60       getAffineForIVs(*memOp, &enclosingOps);
  61       for (unsigned d = 0, e = enclosingOps.size(); d < e; d++) {
  62         AffineForOp loop = enclosingOps[d];
  63         int memRefDim;
  64         bool isContiguous, isInvariant;
  65         if (auto read = dyn_cast<AffineReadOpInterface>(memOp)) {
  66           isContiguous =
  67               isContiguousAccess(loop.getInductionVar(), read, &memRefDim);
  68           isInvariant = isInvariantAccess(read, loop);
  69         } else {
  70           auto write = cast<AffineWriteOpInterface>(memOp);
  71           isContiguous =
  72               isContiguousAccess(loop.getInductionVar(), write, &memRefDim);
  73           isInvariant = isInvariantAccess(write, loop);
  74         }
  75         // Check for contiguity for the innermost memref dimension to avoid
  76         // emitting too many diagnostics.
  77         if (isContiguous && memRefDim == 0)
  78           memOp->emitRemark("contiguous along loop ") << d << '\n';
  79         if (isInvariant)
  80           memOp->emitRemark("invariant along loop ") << d << '\n';
  81       }
  82     }
  83   }
  84 }
  85
  86 namespace mlir {
  87 void registerTestAffineAccessAnalysisPass() {
  88   PassRegistration<TestAccessAnalysis>();
  89 }
  90 } // namespace mlir