Revert "[frontend] Remove Complex Regex for MLIR Parsing (#4924)"

This reverts commit 038cbc5.

Author: whitneywhtsang

python/src/ir.cc

@@ -24,7 +24,6 @@
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Transforms/LocationSnapshot.h"
 
-#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/Types.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
@@ -503,16 +502,6 @@ void init_triton_ir(py::module &&m) {
            [](ModuleOp &self, FuncOp &funcOp) -> void {
              self.push_back(funcOp);
            })
-      .def("get_entry_func_name",
-           [](ModuleOp &self) -> std::string {
-             for (auto &op : self.getOps()) {
-               if (auto func = dyn_cast<FuncOp>(op)) {
-                 if (LLVM::isKernel(func))
-                   return func.getName().str();
-               }
-             }
-             return "";
-           })
       .def("has_function",
            [](ModuleOp &self, std::string &funcName) -> bool {
              if (self.lookupSymbol(funcName))
@@ -523,43 +512,6 @@ void init_triton_ir(py::module &&m) {
            [](ModuleOp &self, std::string &funcName) -> FuncOp {
              return self.lookupSymbol<FuncOp>(funcName);
            })
-      /*
-       * def ty_to_cpp(ty) is the consumer of this function.
-       * If the type is a ptr it expects ty[0] == '*', else the type itself.
-       */
-
-      .def("get_function_signature",
-           [](ModuleOp &self, FuncOp &func) -> std::vector<std::string> {
-             std::vector<std::string> strVec;
-
-             auto type = func.getFunctionType();
-             unsigned numArgs = type.getNumInputs();
-             for (unsigned i = 0; i != numArgs; ++i) {
-               std::string tempType;
-               llvm::raw_string_ostream os(tempType);
-
-               auto ty = type.getInput(i);
-               if (auto attributes = func.getCallableArgAttrs()) {
-                 Attribute attr = attributes[i];
-                 // Check for tt.nv_tma_desc = 1
-                 if (auto dAttr = dyn_cast<DictionaryAttr>(attr)) {
-                   if (dAttr.contains("tt.nv_tma_desc")) {
-                     strVec.push_back("nvTmaDesc");
-                     continue;
-                   }
-                 }
-               }
-               if (auto ptrType = dyn_cast<PointerType>(ty)) {
-                 auto pType = ptrType.getPointeeType();
-                 os << "*";
-                 pType.print(os);
-               } else {
-                 ty.print(os);
-               }
-               strVec.push_back(tempType);
-             }
-             return strVec;
-           })
       .def("get_int_attr",
            [](ModuleOp &self, std::string name) -> py::object {
              auto ret = self->getAttrOfType<IntegerAttr>(name);

python/test/unit/tools/test_irsource.py

@@ -1,7 +1,4 @@
-from .compiler import CompiledKernel, ASTSource, IRSource, compile, make_backend, LazyDict
+from .compiler import CompiledKernel, ASTSource, compile, make_backend, LazyDict
 from .errors import CompilationError
 
-__all__ = [
-    "compile", "make_backend", "ASTSource", "IRSource", "AttrsDescriptor", "CompiledKernel", "CompilationError",
-    "LazyDict"
-]
+__all__ = ["compile", "make_backend", "ASTSource", "AttrsDescriptor", "CompiledKernel", "CompilationError", "LazyDict"]

python/triton/compiler/__init__.py

@@ -25,13 +25,19 @@
 # - (\((?:%\w+: \S+(?: \{\S+ = \S+ : \S+\})?(?:, )?)*\)) : match a pair of parentheses enclosing
 #   zero or more arguments separated by commas, and capture it as group 2 (the argument list)
 # - (attributes \{[\S\s]+\})? : optionally match attributes enclosed in braces and capture it as group 3
+mlir_prototype_pattern = r"^\s*tt\.func\s+(?:public\s+)?(@\w+)(\((?:%\w+: [\S\s]+(?: \{\S+ = \S+ : \S+\})?(?:, )?)*\))\s*(attributes \{[\S\s]+\})?\s+\{\s*$"
 ptx_prototype_pattern = r"\.(?:visible|extern)\s+\.(?:entry|func)\s+(\w+)\s*\(([^)]*)\)"
 prototype_pattern = {
+    "ttir": mlir_prototype_pattern,
+    "ttgir": mlir_prototype_pattern,
     "ptx": ptx_prototype_pattern,
 }
 
+mlir_arg_type_pattern = r'%\w+: ((?:[^,\s<)]+|<[^>]+>)+(?: {[^}]+})?),?'
 ptx_arg_type_pattern = r"\.param\s+\.(\w+)"
 arg_type_pattern = {
+    "ttir": mlir_arg_type_pattern,
+    "ttgir": mlir_arg_type_pattern,
     "ptx": ptx_arg_type_pattern,
 }
 
@@ -49,6 +55,16 @@ def convert_type_repr(x):
     return x
 
 
+def _get_num_warps_from_ir_str(src: str):
+    ttgir_num_warps_pattern = r'"triton_gpu.num-warps"\s?=\s?(\d+)\s?:'
+    # TODO(jlebar): Using a regex to get num-warps is a hack, and will break if
+    # e.g. someone has an instruction (not module) attribute named "num-warps".
+    num_warps_matches = re.findall(ttgir_num_warps_pattern, src)
+    assert len(num_warps_matches) == 1, "Expected exactly one match for num_warps"
+    num_warps = int(num_warps_matches[0])
+    return num_warps
+
+
 class ASTSource:
 
     def __init__(self, fn, signature, constants=None, attrs=None) -> None:
@@ -91,41 +107,28 @@ def parse_options(self):
 
 class IRSource:
 
-    def __init__(self, path, context):
+    def __init__(self, path):
         self.path = path
         path = Path(path)
         self.ext = path.suffix[1:]
         self.src = path.read_text()
-        ir.load_dialects(context)
-
-        # We don't have a easy-to-use PTX parser that we can use, so keep that regex for now.
-        # TODO - replace with a proper parser
-        if self.ext == "ptx":
-            match = re.search(prototype_pattern[self.ext], self.src, re.MULTILINE)
-            self.name = match.group(1)
-            signature = match.group(2)
-            types = re.findall(arg_type_pattern[self.ext], signature)
-            self.signature = {k: convert_type_repr(ty) for k, ty in enumerate(types)}
-        else:
-            self.module = ir.parse_mlir_module(self.path, context)
-            fn_name = self.module.get_entry_func_name()
-            self.name = "@" + fn_name
-            funcOp = self.module.get_function(fn_name)
-            func_ty = self.module.get_function_signature(funcOp)
-            self.signature = {k: ty for k, ty in enumerate(func_ty)}
+        match = re.search(prototype_pattern[self.ext], self.src, re.MULTILINE)
+        self.name = match.group(1)
+        signature = match.group(2)
+        types = re.findall(arg_type_pattern[self.ext], signature)
+        self.signature = {k: convert_type_repr(ty) for k, ty in enumerate(types)}
 
     def hash(self):
         return hashlib.sha256(self.src.encode("utf-8")).hexdigest()
 
     def make_ir(self, options, codegen_fns, module_map, context):
-        self.module.context = context
-        return self.module
+        module = ir.parse_mlir_module(self.path, context)
+        module.context = context
+        return module
 
     def parse_options(self):
         if self.ext == "ttgir":
-            num_warps = self.module.get_int_attr("triton_gpu.num-warps")
-            assert num_warps is not None, "Unable to parse triton_gpu.num-warps attribute"
-            return {'num_warps': num_warps}
+            return {'num_warps': _get_num_warps_from_ir_str(self.src)}
         return dict()
 
 
@@ -222,9 +225,7 @@ def compile(src, target=None, options=None):
     # create backend
     if ir_source:
         assert isinstance(src, str), "source must be either AST or a filepath"
-        context = ir.context()
-        src = IRSource(src, context)
-
+        src = IRSource(src)
     extra_options = src.parse_options()
     options = backend.parse_options(dict(options or dict(), **extra_options))
     # create cache manager
@@ -265,15 +266,9 @@ def compile(src, target=None, options=None):
     # when the source is an IR file, don't apply the passes related to this stage. This makes it easier to write IR level tests.
     if ir_source:
         first_stage += 1
-
-    if not isinstance(src, IRSource):
-        context = ir.context()
-        ir.load_dialects(context)
-        backend.load_dialects(context)
-    else:
-        # For IRSource, we have already grabbed the context + called ir.load_dialects
-        # just need to load the dialects for the backend.
-        backend.load_dialects(context)
+    context = ir.context()
+    ir.load_dialects(context)
+    backend.load_dialects(context)
     codegen_fns = backend.get_codegen_implementation()
     module_map = backend.get_module_map()
     try: