add pybind interface (#9)

1. refactor compiler.py

Author: enjustli

CMakeLists.txt

@@ -23,7 +23,10 @@ add_subdirectory(test)
 add_subdirectory(tools/triton-shared-opt)
 
 if (TRITON_SHARED_BUILD_CPU_BACKEND)
-    add_triton_plugin(TritonShared ${CMAKE_CURRENT_SOURCE_DIR}/triton_shared.cc LINK_LIBS TritonSharedAnalysis TritonTilingExtIR)
+    get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+    get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
+
+    add_triton_plugin(TritonShared ${CMAKE_CURRENT_SOURCE_DIR}/triton_shared.cc LINK_LIBS TritonSharedAnalysis TritonTilingExtIR ${conversion_libs} ${extension_libs} MLIRSparseTensorTransforms MLIRControlFlowTransforms MLIRTensorInferTypeOpInterfaceImpl)
     target_link_libraries(TritonShared PRIVATE Python3::Module pybind11::headers)
 endif()
 

backend/compiler.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 from triton.backends.compiler import BaseBackend, GPUTarget
-from triton._C.libtriton import ir, passes
+from triton._C.libtriton import ir, passes, triton_shared, llvm
 from dataclasses import dataclass
 from typing import Any, Dict, Tuple
 from types import ModuleType
@@ -16,12 +16,6 @@
 import triton
 from pathlib import Path
 
-def _get_triton_shared_opt_path() -> str:
-    path = os.getenv("TRITON_SHARED_OPT_PATH", "")
-    if path == "":
-        raise Exception("TRITON_SHARED_OPT_PATH is not set.")
-    return path
-
 
 def _get_llvm_bin_path(bin_name: str) -> str:
     path = os.getenv("LLVM_BINARY_DIR", "")
@@ -30,13 +24,6 @@ def _get_llvm_bin_path(bin_name: str) -> str:
     return os.path.join(path, bin_name)
 
 
-def _dump_ir_if_needed(files):
-    path = os.getenv("TRITON_SHARED_DUMP_PATH", "")
-    if not path:
-        return
-    for f in files:
-        shutil.copy(f, os.path.join(path, os.path.basename(f)))
-
 def _get_sanitizer_type():
     # returns "" if not set
     # throws error if set to something other than "asan" or "tsan"
@@ -45,96 +32,11 @@ def _get_sanitizer_type():
     if sanitizer_type != "" and sanitizer_type != "asan" and sanitizer_type != "tsan":
         # throw error
         raise Exception(f"TRITON_SHARED_SANITIZER_TYPE {sanitizer_type} is invalid.")
-    
-    return sanitizer_type
 
-def _ttir_to_ttsharedir(mod):
-    # Get Triton-MLIR as string
-    ttir_code = str(mod)
-    with tempfile.TemporaryDirectory() as tmpdir:
-        src_path = os.path.join(tmpdir, "tt.mlir")
-        dst_path = os.path.join(tmpdir, "ttshared.mlir")
-        Path(src_path).write_text(ttir_code)
-        _dump_ir_if_needed([src_path])
-        triton_shared_opt_path = _get_triton_shared_opt_path()
-
-        subprocess_args = [triton_shared_opt_path, src_path, "--triton-to-linalg-experimental", "--mlir-print-debuginfo", "-o", dst_path]
-
-        if _get_sanitizer_type() != "":
-            print("Building with sanitizer support...")
-
-            # has to run before the other passes as operates on the tt dialect
-            subprocess_args.insert(2, "--add-llvm-debug-info")
-
-        subprocess.check_call(subprocess_args)
-        _dump_ir_if_needed([dst_path])
-        return Path(dst_path).read_text()
-
-
-def _optimize_ttsharedir(ttsharedir: str):
-    # We don't apply any optimizations now, but we can add passes if needed.
-    return ttsharedir
+    return sanitizer_type
 
 
-def _ttsharedir_to_llir(ttsharedir: str):
-    with tempfile.TemporaryDirectory() as tmpdir:
-        ttshared_path = os.path.join(tmpdir, "ttshared.mlir")
-        llmlir_path = os.path.join(tmpdir, "ll.mlir")
-        llir_path = os.path.join(tmpdir, "ll.ir")
-        Path(ttshared_path).write_text(ttsharedir)
-        mlir_opt_path = _get_llvm_bin_path("mlir-opt")
-        # TritonShared-MLIR to LLVM-MLIR
-        subprocess.check_call([mlir_opt_path, ttshared_path,
-            "--convert-linalg-to-affine-loops",
-            # Note: eliminate-empty-tensors fails when there are multiple func.return ops
-            # in a single kernel which are the results of early returns.
-            # See python/examples/test_early_return.py for examples.
-            # We disable this pass for now since performance on CPU isn't the main
-            # focus at the moment.
-            # "--eliminate-empty-tensors",
-            "--empty-tensor-to-alloc-tensor",
-            "--one-shot-bufferize=allow-return-allocs-from-loops=true",
-            "--lower-affine",
-            "--convert-linalg-to-loops",
-            "--expand-strided-metadata",
-            "--convert-scf-to-cf",
-            "--convert-arith-to-llvm",
-            "--convert-math-to-llvm",
-            "--convert-complex-to-llvm",
-            "--convert-vector-to-llvm",
-            "--convert-index-to-llvm",
-            "--memref-expand",
-            "--finalize-memref-to-llvm",
-            "--convert-func-to-llvm",
-            "--convert-cf-to-llvm",
-            # Lowering memrefs creates more affine.apply ops.
-            # Lowering these affine ops again creates further arith ops,
-            # so we have to run these two passes again here.
-            "--lower-affine",
-            "--convert-arith-to-llvm",
-            # Remove all unrealized casts created
-            "--reconcile-unrealized-casts",
-            "--mlir-print-debuginfo",
-            "-o",
-            llmlir_path])
-        _dump_ir_if_needed([llmlir_path])
-
-        # LLVM-MLIR to LLVM-IR
-        mlir_translate_path = _get_llvm_bin_path("mlir-translate")
-        subprocess.check_call([mlir_translate_path, llmlir_path,
-            "--mlir-to-llvmir",
-            "-o",
-            llir_path])
-        _dump_ir_if_needed([llir_path])
-        return Path(llir_path).read_text()
-
-
-def _optimize_llir(llir: str):
-    # We don't apply any optimizations now, but we can add passes if needed.
-    return llir
-
-
-def _llir_to_bin(llir: str, metadata):
+def _llir_to_bin(llir: str, metadata, options):
     pattern = r"define void @(\w+)\(.+"
     matches = re.findall(pattern, llir)
     assert len(matches) == 1
@@ -150,17 +52,19 @@ def _llir_to_bin(llir: str, metadata):
             # using a sanitizer
             # invoke pass to append sanitizer attributes
             instrumented_src_path = os.path.join(tmpdir, "kernel-instrumented.ll")
-        
+
             opt_path = _get_llvm_bin_path("opt")
             top_level_triton_path = os.path.dirname(triton.__file__)
-            sanitizer_attributes_pass_path = str(next(Path(top_level_triton_path).rglob("libSanitizerAttributes.so"), None))
+            sanitizer_attributes_pass_path = str(
+                next(Path(top_level_triton_path).rglob("libSanitizerAttributes.so"), None))
 
             if not sanitizer_attributes_pass_path:
                 raise Exception("libSanitizerAttributes.so does not exist.")
 
-            subprocess.check_call([opt_path, "-load-pass-plugin", sanitizer_attributes_pass_path, 
-                "-passes=sanitizer-attributes", f"-sanitizer-type={sanitizer_type}", "-S", src_path, 
-                "-o", instrumented_src_path])
+            subprocess.check_call([
+                opt_path, "-load-pass-plugin", sanitizer_attributes_pass_path, "-passes=sanitizer-attributes",
+                f"-sanitizer-type={sanitizer_type}", "-S", src_path, "-o", instrumented_src_path
+            ])
 
             # compile to object file
             clang_path = _get_llvm_bin_path("clang++")
@@ -171,14 +75,13 @@ def _llir_to_bin(llir: str, metadata):
                 subprocess_args.extend(["-g", "-fsanitize=address", "-mllvm", "-asan-stack=0"])
             elif sanitizer_type == "tsan":
                 subprocess_args.extend(["-g", "-fsanitize=thread"])
-                
+
             subprocess.check_call(subprocess_args)
         else:
             llc_path = _get_llvm_bin_path("llc")
             subprocess.check_call([llc_path, src_path, "-filetype=obj", "-relocation-model=pic", "-o", dst_path])
-        
-        return Path(dst_path).read_bytes()
 
+        return Path(dst_path).read_bytes()
 
 
 @dataclass(frozen=True)
@@ -232,19 +135,13 @@ def pack_metadata(self, metadata):
         # Note: We actually don't need any of these except for the name which is
         # used in the launch function in driver.py. Putting these in so we're
         # consistent with other backends
-        return (
-            metadata.num_warps,
-            metadata.num_ctas,
-            metadata.shared,
-            metadata.cluster_dims[0],
-            metadata.cluster_dims[1],
-            metadata.cluster_dims[2],
-            metadata.name
-        )
+        return (metadata.num_warps, metadata.num_ctas, metadata.shared, metadata.cluster_dims[0],
+                metadata.cluster_dims[1], metadata.cluster_dims[2], metadata.name)
 
     # Our compilation pipeline isn't in python like nvidia or amd, no need to load
     # dialects. See `triton_shared.cc`
     def load_dialects(self, ctx):
+        triton_shared.load_dialects(ctx)
         return
 
     @staticmethod
@@ -263,12 +160,66 @@ def make_ttir(mod, metadata, options):
         pm.run(mod, 'make_ttir')
         return mod
 
+    @staticmethod
+    def make_tt_shared_ir(mod, metadata, options):
+        pm = ir.pass_manager(mod.context)
+        pm.enable_debug()
+        if _get_sanitizer_type() != "":
+            print("Building with sanitizer support...")
+            # has to run before the other passes as operates on the tt dialect
+            triton_shared.add_llvm_debug_info(pm)
+        triton_shared.add_triton_to_linalg_experimental(pm)
+        pm.run(mod, 'make_tt_shared_ir')
+        return mod
+
+    @staticmethod
+    def make_llir(mod, metadata, options):
+        pm = ir.pass_manager(mod.context)
+        pm.enable_debug()
+        triton_shared.add_convert_linalg_to_affine_loops(pm)
+
+        # Note: eliminate-empty-tensors fails when there are multiple func.return ops
+        # in a single kernel which are the results of early returns.
+        # See python/examples/test_early_return.py for examples.
+        # We disable this pass for now since performance on CPU isn't the main
+        # focus at the moment.
+        # triton_shared.add_eliminate_empty_tensors(pm)
+
+        triton_shared.add_empty_tensor_to_alloc_tensor(pm)
+        triton_shared.add_one_shot_bufferize(pm)
+        triton_shared.add_lower_affine(pm)
+        triton_shared.add_convert_linalg_to_loops(pm)
+        triton_shared.add_expand_strided_metadata(pm)
+        triton_shared.add_convert_scf_to_cf(pm)
+        triton_shared.add_convert_arith_to_llvm(pm)
+        triton_shared.add_convert_math_to_llvm(pm)
+        triton_shared.add_convert_complex_to_llvm(pm)
+        triton_shared.add_convert_vector_to_llvm(pm)
+        triton_shared.add_convert_index_to_llvm(pm)
+        triton_shared.add_memref_expand(pm)
+        triton_shared.add_finalize_memref_to_llvm(pm)
+        triton_shared.add_convert_func_to_llvm(pm)
+        triton_shared.add_convert_cf_to_llvm(pm)
+        # Lowering memrefs creates more affine.apply ops.
+        # Lowering these affine ops again creates further arith ops,
+        # so we have to run these two passes again here.
+        triton_shared.add_lower_affine(pm)
+        triton_shared.add_convert_arith_to_llvm(pm)
+        # Remove all unrealized casts created
+        triton_shared.add_reconcile_unrealized_casts(pm)
+        pm.run(mod, 'make_llir')
+
+        # LLVM-IR (MLIR) -> LLVM-IR (LLVM)
+        llvm.init_targets()
+        context = llvm.context()
+        llvm_mod = llvm.to_module(mod, context)
+        return str(llvm_mod)
+
     def add_stages(self, stages, options, language):
         stages["ttir"] = lambda src, metadata: self.make_ttir(src, metadata, options)
-        stages["ttsharedir"] = lambda src, metadata: _optimize_ttsharedir(_ttir_to_ttsharedir(src))
-        stages["llir"] = lambda src, metadata: _optimize_llir(_ttsharedir_to_llir(src))
-        stages["obj"] = lambda src, metadata: _llir_to_bin(src, metadata)
-
+        stages["ttsharedir"] = lambda src, metadata: self.make_tt_shared_ir(src, metadata, options)
+        stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options)
+        stages["obj"] = lambda src, metadata: _llir_to_bin(src, metadata, options)
 
     @functools.lru_cache()
     def hash(self):

python/examples/test_unstructured_mask.py

@@ -7,13 +7,18 @@
 import triton.language as tl
 
 from triton.backends.triton_shared.driver import CPUDriver
-from triton.backends.triton_shared.compiler import _get_triton_shared_opt_path
 
 import os
 from pathlib import Path
 import subprocess
 import tempfile
 
+def _get_triton_shared_opt_path() -> str:
+    path = os.getenv("TRITON_SHARED_OPT_PATH", "")
+    if path == "":
+        raise Exception("TRITON_SHARED_OPT_PATH is not set.")
+    return path
+
 def run_triton_to_structured(ttir_code):
     with tempfile.TemporaryDirectory() as tmpdir:
         src_path = os.path.join(tmpdir, "tt.mlir")