Change to analysis implementation

Author: Alwaysproblem

mlir/optimization/scheduler/lib/OpStatsPass.cpp

@@ -2,17 +2,29 @@
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Pass/Pass.h"
+#include <cstdint>
 
 using namespace mlir;
 
 namespace {
-struct LabOpStatsPass
-    : public PassWrapper<LabOpStatsPass, OperationPass<func::FuncOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LabOpStatsPass)
 
-  void runOnOperation() override {
-    func::FuncOp func = getOperation();
+class LabOpStateAnalysis {
+public:
+  struct States {
+    int64_t M;
+    int64_t N;
+    int64_t K;
+    int64_t flops;
+    int64_t bytes;
+    int64_t macs;
+    double intensity;
+    unsigned numLoops;
+    unsigned numParallel;
+    unsigned numReduction;
+  };
 
+  explicit LabOpStateAnalysis(Operation *op) {
+    auto func = cast<func::FuncOp>(op);
     func.walk([&](Operation *op) {
       if (auto matmul = dyn_cast<linalg::MatmulOp>(op)) {
         analyzeMatmul(matmul);
@@ -54,6 +66,7 @@ struct LabOpStatsPass
     }
 
     int64_t flops = 2 * M * N * K;
+    int64_t totalMacs = M * N * K;
     int64_t aBytes = aType.getNumElements() * elemBytes;
     int64_t bBytes = bType.getNumElements() * elemBytes;
     int64_t cBytes = cType.getNumElements() * elemBytes;
@@ -63,20 +76,49 @@ struct LabOpStatsPass
                                             static_cast<double>(totalBytes)
                                       : 0.0;
 
-    op.emitRemark() << "[lab-op-stats] matmul "
-                    << "M=" << M << " N=" << N << " K=" << K
-                    << " flops=" << flops << " bytes=" << totalBytes
-                    << " intensity=" << intensity;
+    states.M = M;
+    states.N = N;
+    states.K = K;
+    states.flops = flops;
+    states.bytes = totalBytes;
+    states.macs = totalMacs;
+    states.intensity = intensity;
   }
 
   void analyzeGeneric(linalg::GenericOp op) {
     unsigned numLoops = op.getNumLoops();
     unsigned numParallel = op.getNumParallelLoops();
     unsigned numReduction = numLoops - numParallel;
 
-    op.emitRemark() << "[lab-op-stats] generic "
-                    << "loops=" << numLoops << " parallel=" << numParallel
-                    << " reduction=" << numReduction;
+    states.numLoops = numLoops;
+    states.numParallel = numParallel;
+    states.numReduction = numReduction;
+  }
+
+  const States &getStates() const { return states; }
+
+private:
+  States states;
+};
+
+struct LabOpStatsPass
+    : public PassWrapper<LabOpStatsPass, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LabOpStatsPass)
+
+  StringRef getArgument() const override { return "lab-op-stats"; }
+
+  void runOnOperation() override {
+    auto func = getOperation();
+    auto &analysis = getAnalysis<LabOpStateAnalysis>();
+    const auto &states = analysis.getStates();
+
+    func.emitRemark() << "[lab-op-stats] M=" << states.M << " N=" << states.N
+                      << " K=" << states.K << " FLOPs=" << states.flops
+                      << " Bytes=" << states.bytes << " MACs=" << states.macs
+                      << " Intensity=" << states.intensity
+                      << " Loops=" << states.numLoops
+                      << " Parallel=" << states.numParallel
+                      << " Reduction=" << states.numReduction;
   }
 };
 } // namespace