Introducing SCF

include/triton/ir/Dialect.h

@@ -6,6 +6,7 @@
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/SCF/SCF.h"
 // #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 
 #include "triton/ir/Dialect.h.inc"

lib/ir/CMakeLists.txt

@@ -17,4 +17,5 @@ add_mlir_dialect_library(TritonIR
   MLIRStandard
 
   MLIRTensor
+  MLIRSCF
 )

python/src/triton.cc

@@ -834,13 +834,17 @@ void init_triton_ir(py::module &&m) {
         }
         throw std::runtime_error("invalid function type");
       })
-      // // Structured control flow
-      // .def("create_scf_for", [](mlir::OpBuilder &self) {
-      //   return self.create<mlir::scf::ForOp>(/*fill this*/);
-      // })
-      // .def("create_scf_yield")
-      // .def("create_scf_if")
-      // .def("create_scf_while")
+      // Structured control flow
+      .def("create_for", [](mlir::OpBuilder &self, MlirValue &lb, MlirValue &ub,
+                            MlirValue &step) {
+        auto loc = self.getUnknownLoc();
+        return wrap(
+          self.create<mlir::scf::ForOp>(loc, unwrap(lb), unwrap(ub), unwrap(step))
+        );
+      })
+      // .def("create_yield")
+      // .def("create_if")
+      // .def("create_while")
 
       // miscellious
       .def("create_make_range", [](mlir::OpBuilder &self, int start, int end) -> MlirValue {