68 lines
1.8 KiB
C++
68 lines
1.8 KiB
C++
#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>(); }
|
|
}
|
|
}
|
|
|