Removed get_*ir_passes()

Author: AndreyPavlenko

third_party/intel/backend/compiler.py

@@ -127,6 +127,10 @@ def parse_target(self, tgt_prop) -> dict:
         dev_prop['has_subgroup_2d_block_io'] = tgt_prop.get('has_subgroup_2d_block_io', False)
         dev_prop['has_bfloat16_conversions'] = tgt_prop.get('has_bfloat16_conversions', True)
 
+        if self.device_arch in self.device_props:
+            dev_prop.update(self.device_props[self.device_arch])
+            return dev_prop
+
         return dev_prop
 
     def parse_options(self, opts) -> Any:
@@ -202,85 +206,27 @@ def get_split_barrier_scope(opt):
             split_barriers_scope = intel.SplitBarrierScope.Subgroup
         return split_barriers_scope
 
-    @classmethod
-    def create_pass_manager(cls, context, add_passes=[]):
-        pm = ir.pass_manager(context)
-        pm.enable_debug()
-        for p in add_passes:
-            if p is None:
-                continue
-            elif isinstance(p, tuple):
-                p[0](pm, *p[1:])
-            else:
-                p(pm)
-        return pm
-
-    @classmethod
-    def get_ttir_passes(cls, opt):
-        return [
-            passes.common.add_inliner,
-            intel.passes.ttir.add_convert_tdesc_to_block_pointer,
-            passes.ttir.add_rewrite_tensor_descriptor_to_pointer,
-            passes.common.add_cse,
-            passes.common.add_licm,
-            intel.passes.ttir.add_remove_masks,
-            intel.passes.ttir.add_fuse_reshape,
-            passes.common.add_canonicalizer,
-            passes.ttir.add_combine,
-            passes.ttir.add_reorder_broadcast,
-            passes.common.add_cse,
-            passes.common.add_symbol_dce,
-            passes.ttir.add_loop_unroll,
-        ]
-
     @classmethod
     @track
     def make_ttir(cls, mod, metadata, opt):
-        pm = cls.create_pass_manager(mod.context, cls.get_ttir_passes(opt))
+        pm = ir.pass_manager(mod.context)
+        pm.enable_debug()
+        passes.common.add_inliner(pm)
+        intel.passes.ttir.add_convert_tdesc_to_block_pointer(pm)
+        passes.ttir.add_rewrite_tensor_descriptor_to_pointer(pm)
+        passes.common.add_cse(pm)
+        passes.common.add_licm(pm)
+        intel.passes.ttir.add_remove_masks(pm)
+        intel.passes.ttir.add_fuse_reshape(pm)
+        passes.common.add_canonicalizer(pm)
+        passes.ttir.add_combine(pm)
+        passes.ttir.add_reorder_broadcast(pm)
+        passes.common.add_cse(pm)
+        passes.common.add_symbol_dce(pm)
+        passes.ttir.add_loop_unroll(pm)
         pm.run(mod, 'make_ttir')
         return mod
 
-    @classmethod
-    def get_ttgir_passes(cls, opt):
-        # fmt: off
-        return [
-            (passes.ttir.add_convert_to_ttgpuir, "xpu", opt.num_warps, opt.warp_size, opt.num_ctas),
-            # optimize TTGIR
-            intel.passes.ttgpuir.add_coalesce,
-            intel.passes.ttgpuir.add_remove_layout_conversions,
-
-            intel.passes.ttgpuir.add_accelerate_matmul,
-            intel.passes.ttgpuir.add_materialize_block_pointer,
-            intel.passes.ttgpuir.add_remove_layout_conversions,
-            intel.passes.ttgpuir.add_optimize_dot_operands,
-            (intel.passes.ttgpuir.add_pipeline, opt.num_stages, cls.get_split_barrier_scope(opt)),
-
-            intel.passes.ttgpuir.add_reduce_variable_liveness if opt.reduce_variable_liveness else None,
-
-            passes.ttgpuir.add_fuse_nested_loops,
-
-            passes.common.add_canonicalizer,
-            passes.ttir.add_triton_licm,
-            passes.common.add_canonicalizer,
-            passes.ttgpuir.add_combine_tensor_select_and_if,
-
-            passes.ttgpuir.add_optimize_thread_locality,
-            (passes.ttgpuir.add_optimize_dot_operands, True),
-            passes.common.add_cse,
-            passes.ttgpuir.add_prefetch,
-            (passes.ttgpuir.add_optimize_dot_operands, True),
-            intel.passes.ttgpuir.add_remove_layout_conversions,
-            intel.passes.ttgpuir.add_reduce_data_duplication,
-            passes.ttgpuir.add_reorder_instructions,
-            passes.common.add_cse,
-            passes.common.add_symbol_dce,
-            passes.common.add_sccp,
-            passes.common.add_canonicalizer,
-            intel.passes.ttgpuir.add_optimize_reduction_locality if knobs.intel.opt_reduction_locality else None,
-            (intel.passes.arith.add_arith_emulate_unsupported_floats, ["bf16"], "f32")
-        ]
-        # fmt: on
-
     @classmethod
     @track
     def make_ttgir(cls, mod, metadata, opt, properties):
@@ -291,7 +237,8 @@ def make_ttgir(cls, mod, metadata, opt, properties):
             cluster_info.clusterDimZ = opt.cluster_dims[2]
 
         # Annotate module with information required by subsequent transformations.
-        pm = cls.create_pass_manager(mod.context)
+        pm = ir.pass_manager(mod.context)
+        pm.enable_debug()
         module_opts = intel.passes.ttgpuir.AnnotateModuleOptions()
         cls.annotate_module(module_opts, properties, opt)
         intel.passes.ttgpuir.add_triton_annotate_module(pm, module_opts)
@@ -301,7 +248,44 @@ def make_ttgir(cls, mod, metadata, opt, properties):
         opt.warp_size = intel.get_threads_per_warp(mod)
         cls.validate_options(opt, properties)
 
-        pm = cls.create_pass_manager(mod.context, cls.get_ttgir_passes(opt))
+        pm = ir.pass_manager(mod.context)
+        pm.enable_debug()
+        passes.ttir.add_convert_to_ttgpuir(pm, "xpu", opt.num_warps, opt.warp_size, opt.num_ctas)
+        # optimize TTGIR
+        intel.passes.ttgpuir.add_coalesce(pm)
+        intel.passes.ttgpuir.add_remove_layout_conversions(pm)
+
+        intel.passes.ttgpuir.add_accelerate_matmul(pm)
+        intel.passes.ttgpuir.add_materialize_block_pointer(pm)
+        intel.passes.ttgpuir.add_remove_layout_conversions(pm)
+        intel.passes.ttgpuir.add_optimize_dot_operands(pm)
+        intel.passes.ttgpuir.add_pipeline(pm, opt.num_stages, XPUBackend.get_split_barrier_scope(opt))
+
+        if (opt.reduce_variable_liveness):
+            intel.passes.ttgpuir.add_reduce_variable_liveness(pm)
+
+        passes.ttgpuir.add_fuse_nested_loops(pm)
+
+        passes.common.add_canonicalizer(pm)
+        passes.ttir.add_triton_licm(pm)
+        passes.common.add_canonicalizer(pm)
+        passes.ttgpuir.add_combine_tensor_select_and_if(pm)
+
+        passes.ttgpuir.add_optimize_thread_locality(pm)
+        passes.ttgpuir.add_optimize_dot_operands(pm, True)
+        passes.common.add_cse(pm)
+        passes.ttgpuir.add_prefetch(pm)
+        passes.ttgpuir.add_optimize_dot_operands(pm, True)
+        intel.passes.ttgpuir.add_remove_layout_conversions(pm)
+        intel.passes.ttgpuir.add_reduce_data_duplication(pm)
+        passes.ttgpuir.add_reorder_instructions(pm)
+        passes.common.add_cse(pm)
+        passes.common.add_symbol_dce(pm)
+        passes.common.add_sccp(pm)
+        passes.common.add_canonicalizer(pm)
+        if knobs.intel.opt_reduction_locality:
+            intel.passes.ttgpuir.add_optimize_reduction_locality(pm)
+        intel.passes.arith.add_arith_emulate_unsupported_floats(pm, ["bf16"], "f32")
         pm.run(mod, 'make_ttgir')
         metadata["cluster_dims"] = (cluster_info.clusterDimX, cluster_info.clusterDimY, cluster_info.clusterDimZ)
         return mod
@@ -322,31 +306,6 @@ def gluon_to_ttgir(self, src, metadata, options):
         metadata["tensordesc_meta"] = mod.get_tensordesc_metadata()
         return mod
 
-    @classmethod
-    def get_llir_passes(cls, opt, mod):
-        # fmt: off
-        return [
-            passes.convert.add_scf_to_cf,
-            passes.gluon.add_inliner,
-            passes.convert.add_index_to_llvmir,
-            intel.passes.ttgpuir.add_allocate_shared_memory,
-            passes.ttgpuir.add_allocate_global_scratch_memory,
-            # instrumentation point here so we can override IRs above (e.g., ttir and ttgir)
-            lambda pm: cls.instrumentation.patch("ttgpuir_to_llvmir", pm, mod.context) if cls.instrumentation else None,
-            intel.passes.ttgpuir.add_to_llvmir,
-            intel.passes.ttgpuir.add_gen_to_llvm,
-            passes.common.add_canonicalizer,
-            intel.passes.ttgpuir.add_rewrite_stack_ptr,
-            passes.common.add_cse,
-            passes.convert.add_arith_to_llvmir,
-            passes.common.add_canonicalizer,
-            passes.common.add_cse,
-            passes.common.add_symbol_dce,
-            None if knobs.compilation.disable_line_info or knobs.compilation.dump_ir_extract_di_local_variables else passes.llvmir.add_di_scope,
-            lambda pm: cls.instrumentation.patch("llvmir_to_llvm", pm, mod.context) if cls.instrumentation else None,
-        ]
-        # fmt: on
-
     @classmethod
     def optimize_llvm_mod(cls, llvm_mod, options):
         intel.set_spv_target_triple(llvm_mod)
@@ -358,21 +317,50 @@ def optimize_llvm_mod(cls, llvm_mod, options):
     def make_llir(cls, src, metadata, options):
         mod = src
         # TritonGPU -> LLVM-IR (MLIR)
-        pm = cls.create_pass_manager(mod.context, cls.get_llir_passes(options, mod))
+        pm = ir.pass_manager(mod.context)
+        pm.enable_debug()
+
+        passes.convert.add_scf_to_cf(pm)
+        passes.gluon.add_inliner(pm)
+        passes.convert.add_index_to_llvmir(pm)
+        intel.passes.ttgpuir.add_allocate_shared_memory(pm)
+        passes.ttgpuir.add_allocate_global_scratch_memory(pm)
+        # instrumentation point here so we can override IRs above (e.g., ttir and ttgir)
+        if cls.instrumentation:
+            cls.instrumentation.patch("ttgpuir_to_llvmir", pm, mod.context)
+        intel.passes.ttgpuir.add_to_llvmir(pm)
+        intel.passes.ttgpuir.add_gen_to_llvm(pm)
+        passes.common.add_canonicalizer(pm)
+        intel.passes.ttgpuir.add_rewrite_stack_ptr(pm)
+        passes.common.add_cse(pm)
+        passes.convert.add_arith_to_llvmir(pm)
+        passes.common.add_canonicalizer(pm)
+        passes.common.add_cse(pm)
+        passes.common.add_symbol_dce(pm)
+
+        if not knobs.compilation.disable_line_info and not knobs.compilation.dump_ir_extract_di_local_variables:
+            passes.llvmir.add_di_scope(pm)
+
+        if cls.instrumentation:
+            cls.instrumentation.patch("llvmir_to_llvm", pm, mod.context)
         pm.run(mod, 'make_llir')
 
         if knobs.compilation.dump_ir_extract_di_local_variables:
             # comments below on why separate it
             if not knobs.compilation.disable_line_info:
-                pm = cls.create_pass_manager(mod.context, [passes.llvmir.add_di_scope])
+                pm = ir.pass_manager(mod.context)
+                pm.enable_debug()
+                passes.llvmir.add_di_scope(pm)
                 pm.run(mod, 'make_llir.disable_line_info')
 
             # insert dbg intrinsic with several DI Attribute including source
             # var name and type info note: unknown reason for now, but this
             # pass and add_di_scope has to be run separately, otherwise if we
             # put them into previous pipline, it trigger a segmentfault without
             # any error message; could be due to a bug in mlir or pybind11
-            pm = cls.create_pass_manager(mod.context, [passes.llvmir.add_di_local_variable])
+            pm = ir.pass_manager(mod.context)
+            pm.enable_debug()
+            passes.llvmir.add_di_local_variable(pm)
             pm.run(mod, 'make_llir.dump_ir_extract_di_local_variables')
 
         # LLVM-IR (MLIR) -> LLVM-IR (LLVM)