[Analysis] Added Axis Info Analysis (#8)

test/lib/Analysis/CMakeLists.txt

@@ -0,0 +1,6 @@
+add_mlir_library(TritonTestAnalysis
+  TestAxisInfo.cpp
+
+  LINK_LIBS PUBLIC
+  TritonAnalysis
+)

test/lib/Analysis/TestAxisInfo.cpp

@@ -0,0 +1,67 @@
+#include "triton/Analysis/AxisInfo.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace{
+
+struct TestAxisInfoPass
+    : public PassWrapper<TestAxisInfoPass,  OperationPass<FuncOp>>{
+      
+  // LLVM15+
+  // MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAlignmentPass);
+  
+  void print(const std::string& name, raw_ostream& os, ArrayRef<int> vals){
+    os << name << ": [";
+    for(size_t d = 0; d < vals.size(); d++){
+      if(d != 0) os << ", ";
+      os << vals[d];
+    }
+    os << "]";
+  }
+
+  StringRef getArgument() const final { return "test-print-alignment"; }
+  StringRef getDescription() const final 
+  { return "print the result of the alignment analysis pass"; }
+
+  void runOnOperation() override {
+    Operation* operation = getOperation();
+    auto& os = llvm::errs();
+    os << "Testing: " << operation->getName() << "\n";
+    AxisInfoAnalysis analysis(&getContext());
+    analysis.run(operation);
+    operation->walk([&](Operation* op){
+      if(op->getNumResults() < 1)
+        return;
+      for(Value result: op->getResults()){
+        // std::ostringstream oss;
+        // result.print(oss);
+        // os << " => ";
+        LatticeElement<AxisInfo> *latticeElement = analysis.lookupLatticeElement(result);
+        if(!latticeElement){
+          os << "None\n";
+          return;
+        }
+        AxisInfo& info = latticeElement->getValue();
+        print("Contiguity", os, info.getContiguity());
+        os << " ; ";
+        print("Divisibility", os, info.getDivisibility());
+        os << " ; ";
+        print("Constancy", os, info.getConstancy());
+        os << " ( ";
+        result.print(os);
+        os << " ) "; 
+        os << "\n";
+      }
+    });
+  }
+};
+
+}
+
+namespace mlir{
+namespace test{
+void registerTestAlignmentPass() { PassRegistration<TestAxisInfoPass>(); }
+}
+}
+

test/lib/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(Analysis)