gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Analysis] Added Axis Info Analysis (#8)

Browse Source
This commit is contained in:
Philippe Tillet
2022-07-19 13:38:48 -07:00
committed by GitHub
parent df940aaab0
commit a633d2b403
20 changed files with 582 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/src/triton.cc
View File

@@ -788,6 +788,11 @@ void init_triton_ir(py::module &&m) {
.def("add_entry_block", [](mlir::FuncOp &self) -> mlir::Block* {
return self.addEntryBlock();
}, ret::reference)
.def("set_arg_attr", [](mlir::FuncOp &self, int arg_no, const std::string& name, int val){
// set arg attributes "name" to value "val"
auto attrTy = mlir::IntegerType::get(self.getContext(), 32);
self.setArgAttr(arg_no, name, mlir::IntegerAttr::get(attrTy, val));
}, ret::reference)
.def("reset_type", &mlir::FuncOp::setType)
;
@@ -1265,9 +1270,10 @@ void init_triton_ir(py::module &&m) {
.def("create_splat", [](mlir::OpBuilder &self, mlir::Value &arg, std::vector<int64_t> &shape) -> mlir::Value {
auto loc = self.getUnknownLoc();
auto argType = arg.getType();
return self.create<mlir::triton::BroadcastOp>(
auto ret = self.createOrFold<mlir::triton::SplatOp>(
loc, mlir::RankedTensorType::get(shape, argType), arg
);
return ret;
})
// // atomic
.def("create_atomic_cas", [](mlir::OpBuilder &self, mlir::Value &ptr,
@@ -1337,6 +1343,12 @@ void init_triton_ir(py::module &&m) {
.def("run", [](mlir::PassManager &self, mlir::ModuleOp &mod) -> bool {
return mlir::succeeded(self.run(mod.getOperation()));
})
.def("add_sccp_pass", [](mlir::PassManager &self) {
self.addPass(mlir::createSCCPPass());
})
.def("add_symbol_dce_pass", [](mlir::PassManager &self) {
self.addPass(mlir::createSymbolDCEPass());
})
.def("add_inliner_pass", [](mlir::PassManager &self) {
self.addPass(mlir::createInlinerPass());
})