Merge pull request #167 from PSAL-POSTECH/autotune

[Frontend] Template autotune

Author: YWHyuk

PyTorchSimFrontend/extension_codecache.py

@@ -299,7 +299,7 @@ def dummy_simulator(*args, **kwargs):
                 # Dump arguments and meta data
                 dump_metadata(args, arg_attributes, result_path)
                 runtime_path = FunctionalSimulator.get_runtime_dump_path(result_path)
-                if extension_config.CONFIG_TORCHSIM_VALIDATION_MODE or validate:
+                if not autotune and (extension_config.CONFIG_TORCHSIM_VALIDATION_MODE or validate):
                     funcsim = FunctionalSimulator(result_path, key)
                     funcsim.run_spike(args, arg_attributes,
                                     runtime_path, self.validation_binary_name,

PyTorchSimFrontend/extension_config.py

@@ -46,7 +46,9 @@
 
 # AUTOTUNE config
 CONFIG_AUTOTUNE = int(os.environ.get('AUTOTUNE', default=True))
+CONFIG_AUTOTUNE_TEMPLATE = int(os.environ.get('AUTOTUNE_TEMPLATE', default=True))
 CONFIG_MAX_AUTOTUNE_TRY = int(os.environ.get('MAX_AUTOTUNE_TRY', default=10))
+CONFIG_AUTOTUNE_TEMPLATE_TOPK = int(os.environ.get('AUTOTUNE_TEMPLATE_TOPK', default=4))
 
 # For block sparse
 CONFIG_BLOCK_SPARSE = int(os.environ.get('BLOCK_SPARSE', default=0))

PyTorchSimFrontend/mlir/mlir_autotune.py

@@ -74,7 +74,8 @@ def cached_run_fn(*args, **kwargs):
             self.source_code, vectorlane_size=self.extra_args["vector_lane"],
             loop_size=None, spad_info=self.extra_args["spad_info"],
             vlen=self.extra_args["vlen"], arg_attributes=self.extra_args["arg_attributes"],
-            origins="Unknown", silent_mode=True)
+            origins="Unknown", silent_mode=True,
+            validate=self.extra_args['validate'], autotune=self.extra_args['autotune'])
 
         args = [
             tensor

PyTorchSimFrontend/mlir/mlir_bmm_template.py

@@ -6,8 +6,6 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplate
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
-from torch._inductor.codecache import write_atomic
-import PyTorchSimFrontend.extension_codecache as extension_codecache
 from PyTorchSimFrontend.mlir import mlir_common
 
 BMM_TEMPLATE = r"""
@@ -162,51 +160,31 @@ def render(self,
                template_buffer_node = None,
                epilogue_nodes: Optional[List[IRNode]] = None,
                prologue_nodes: Optional[List[IRNode]] = None,
+               tile_info = None,
                **kwargs):
-        if template_buffer_node is not None:
-            self.output_node = template_buffer_node
-
-        # Extract input arguments info
-        X, W = self.input_nodes[0], self.input_nodes[1]
-        Y = self.output_node
-        Bias = None if len(self.input_nodes) == 2 else self.input_nodes[2]
-
-        W_tensor =  empty_strided(W.layout.size, W.layout.stride)
-        X_tensor =  empty_strided(X.layout.size, X.layout.stride)
-        if len(W_tensor.size()) > 3 or len(W_tensor.size()) == 2:
-          W_tensor = W_tensor.view([-1, W_tensor.shape[-2], W_tensor.shape[-1]])
-        if len(X_tensor.size()) > 3 or len(X_tensor.size()) == 2:
-          X_tensor = X_tensor.view([-1, X_tensor.shape[-2], X_tensor.shape[-1]])
-        B, M, N, K = X_tensor.size()[0], X_tensor.size()[1], W_tensor.size()[2], X_tensor.size()[2]
-
-        W_stride = W_tensor.stride()
-        X_stride = X_tensor.stride()
-
-        # Select tile size
-        n_extra_node = len(epilogue_nodes) if epilogue_nodes is not None else 0
-        TILE_M, TILE_N, TILE_K = kernel.gemm_combination_mapping(M, N, K, n_extra_node=n_extra_node)
-        SUB_TILE_M = TILE_M if (TILE_M < kernel.vector_lane) or prologue_nodes else kernel.vector_lane
-        SUB_TILE_N = TILE_N # if (TILE_N < kernel.vector_lane) or prologue_nodes else kernel.vector_lane
-        SUB_TILE_K = TILE_K # if (TILE_K < kernel.vector_lane) or prologue_nodes else kernel.vector_lane
+        X, W, Y, Bias, W_tensor, X_tensor, B, M, N, K, n_extra_node, n_prologue_node = self.extract_info(template_buffer_node, epilogue_nodes, prologue_nodes)
+        if tile_info is None:
+            TILE_M, TILE_N, TILE_K, SUB_TILE_M, SUB_TILE_N, SUB_TILE_K = self.select_tile(kernel, M, N, K, n_extra_node, 0, n_prologue_node)[0]
+        else:
+            TILE_M, TILE_N, TILE_K, SUB_TILE_M, SUB_TILE_N, SUB_TILE_K = tile_info
 
         TOG_latency = M if TILE_M > M else TILE_M
         kernel.loop_size = [TOG_latency, TILE_N, TILE_K]
-        TILE_K = TILE_K // 2 if prologue_nodes else TILE_K
 
         # Select template code
         nr_reduction_nodes = [node for node in epilogue_nodes if node.is_reduction()] if epilogue_nodes is not None else []
         if nr_reduction_nodes:
-          template = BMM_REDUCTION_TEMPLATE
-          epilogue_dim_aliasing = {"index0":"index0", "index1":"index2", "index2": "index1"}
-          nr_rdim = 1
+            template = BMM_REDUCTION_TEMPLATE
+            epilogue_dim_aliasing = {"index0":"index0", "index1":"index2", "index2": "index1"}
+            nr_rdim = 1
         elif prologue_nodes:
-          template = BMM_PROLOGUE_TEMPLATE
-          epilogue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2": "index2"}
-          nr_rdim = 0
+            template = BMM_PROLOGUE_TEMPLATE
+            epilogue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2": "index2"}
+            nr_rdim = 0
         else:
-          template = BMM_TEMPLATE
-          epilogue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2": "index2"}
-          nr_rdim = 0
+            template = BMM_TEMPLATE
+            epilogue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2": "index2"}
+            nr_rdim = 0
 
         # Prepare tile descriptors
         vlane_stride = 1
@@ -323,19 +301,53 @@ def render(self,
             dram_idx = Y_idx,
             dram_tile_desc = Y_tile_desc,
             nr_rdim = nr_rdim,
+            r_dim_size = M,
             dim_aliasing = epilogue_dim_aliasing
         )
         code = self._template_from_string(template).render(**kernel.render_options)
         kernel.add_loop_info([kernel.render_options["M"], kernel.render_options["N"], kernel.render_options["K"]], [kernel.render_options["TILE_M"], kernel.render_options["TILE_N"], kernel.render_options["TILE_K"]])
         return code
 
-    def codegen_header(self, code, extra_headers):
-        write_path = extension_codecache.get_write_path(code)
-        if not os.path.exists(write_path):
-            os.makedirs(write_path)
-        spike_write_path = os.path.join(write_path, "global_var.h")
-        gem5_write_path = os.path.join(write_path, "gem5_global_var.h")
-        if not os.path.exists(spike_write_path):
-            write_atomic(spike_write_path, extra_headers[0])
-        if not os.path.exists(gem5_write_path):
-            write_atomic(gem5_write_path, extra_headers[1])
+    def extract_info(self, template_buffer_node, epilogue_nodes, prologue_nodes):
+        if template_buffer_node is not None:
+            self.output_node = template_buffer_node
+
+        # Extract input arguments info
+        X, W = self.input_nodes[0], self.input_nodes[1]
+        Y = self.output_node
+        Bias = None if len(self.input_nodes) == 2 else self.input_nodes[2]
+
+        W_tensor =  empty_strided(W.layout.size, W.layout.stride)
+        X_tensor =  empty_strided(X.layout.size, X.layout.stride)
+        if len(W_tensor.size()) > 3 or len(W_tensor.size()) == 2:
+          W_tensor = W_tensor.view([-1, W_tensor.shape[-2], W_tensor.shape[-1]])
+        if len(X_tensor.size()) > 3 or len(X_tensor.size()) == 2:
+          X_tensor = X_tensor.view([-1, X_tensor.shape[-2], X_tensor.shape[-1]])
+        B, M, N, K = X_tensor.size()[0], X_tensor.size()[1], W_tensor.size()[2], X_tensor.size()[2]
+
+        W_stride = W_tensor.stride()
+        X_stride = X_tensor.stride()
+
+        # Select tile size
+        n_extra_node = len(epilogue_nodes) if epilogue_nodes is not None else 0
+        n_prologue_node = len(prologue_nodes) if prologue_nodes is not None else 0
+        return X,W,Y,Bias,W_tensor,X_tensor,B,M,N,K,n_extra_node, n_prologue_node
+
+    def get_tile_candidates(self,
+               kernel: MLIRTemplateKernel,
+               template_buffer_node = None,
+               epilogue_nodes: Optional[List[IRNode]] = None,
+               prologue_nodes: Optional[List[IRNode]] = None,
+               **kwargs):
+        X, W, Y, Bias, W_tensor, X_tensor, B, M, N, K, n_extra_node, n_prologue_node = self.extract_info(template_buffer_node, epilogue_nodes, prologue_nodes)
+        return self.select_tile(kernel, M, N, K, n_extra_node, 0, n_prologue_node)
+
+    def select_tile(self, kernel, M, N, K, n_extra_node, n_extra_read, n_prologue_node):
+        tile_candidates = kernel.gemm_combination_mapping(M, N, K, n_extra_node=n_extra_node)
+        for idx, (TILE_M, TILE_N, TILE_K) in enumerate(tile_candidates):
+            SUB_TILE_M = TILE_M if (TILE_M < kernel.vector_lane) or n_prologue_node else kernel.vector_lane
+            SUB_TILE_N = TILE_N # if (TILE_N < kernel.vector_lane) or prologue_nodes else kernel.vector_lane
+            SUB_TILE_K = TILE_K # if (TILE_K < kernel.vector_lane) or prologue_nodes else kernel.vector_lane
+            TILE_K = TILE_K // 2 if n_prologue_node else TILE_K
+            tile_candidates[idx] = TILE_M,TILE_N,TILE_K,SUB_TILE_M,SUB_TILE_N,SUB_TILE_K
+        return tile_candidates

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -17,8 +17,7 @@
     sympy_product
 )
 from torch.utils._sympy.functions import ModularIndexing, FloorDiv
-import PyTorchSimFrontend.extension_codecache as extension_codecache
-
+from PyTorchSimFrontend import extension_codecache
 from PyTorchSimFrontend import extension_config
 from . import mlir_common
 from .mlir_common import LoopLevel, LoopNest
@@ -1565,10 +1564,10 @@ def make_choices(self, nodes, kernel_name):
             current_tile_sz = tuple(self.kernel_group.tile_desc.get_tile_size())
             search_space.add(current_tile_sz)
 
-            print(f"[Auto-tune] Trying tile size: {current_tile_sz}, vlane_stride: {vlane_stride}, split_axis: {self.kernel_group.tile_desc.vmap.vlane_split_axis}")
+            print(f"[Auto-tune] Trying tile size: {list(current_tile_sz)}, vlane_stride: {self.kernel_group.tile_desc.vmap.vlane_stride}, split_axis: {self.kernel_group.tile_desc.vmap.vlane_split_axis}")
             self._prepare_simulator_headers(src_code)
             bench_runner = self.run_bench(nodes, kernel_name, src_code)
-            choices.append((bench_runner, src_code, self.kernel_group))
+            choices.append((bench_runner, src_code, current_tile_sz, self.kernel_group.tile_desc.vmap.vlane_stride))
 
             while prevent_infinite_loop < 10 and candidate_axes:
                 for axis in list(candidate_axes):
@@ -1593,33 +1592,39 @@ def make_choices(self, nodes, kernel_name):
                     src_code = super().codegen_nodes(nodes, kernel_name)
                     current_tile_sz = tuple(self.kernel_group.tile_desc.get_tile_size())
 
+                    # FIXME. How to intergrate this constraint to tile system?
+                    pad = self.kernel_group.tile_desc.vmap.get_used_vlane(current_tile_sz) * self.kernel_group.tile_desc.vmap.vlane_stride
+                    vlane_size = current_tile_sz[self.kernel_group.tile_desc.vmap.vlane_split_axis]
+                    if vlane_size > pad and vlane_size % pad:
+                        prevent_infinite_loop += 1
+                        continue
+
                     # If tile size is converged for this axis, remove from candidate axes
                     if current_tile_sz in search_space:
                         candidate_axes.remove(axis)
                         continue
 
                     # Add this choice
                     search_space.add(current_tile_sz)
-                    print(f"[Auto-tune] Trying tile size: {current_tile_sz}, vlane_stride: {vlane_stride}, split_axis: {self.kernel_group.tile_desc.vmap.vlane_split_axis}")
+                    print(f"[Auto-tune] Trying tile size: {list(current_tile_sz)}, vlane_stride: {self.kernel_group.tile_desc.vmap.vlane_stride}, split_axis: {self.kernel_group.tile_desc.vmap.vlane_split_axis}")
                     self._prepare_simulator_headers(src_code)
                     bench_runner = self.run_bench(nodes, kernel_name, src_code)
-                    choices.append((bench_runner, src_code, self.kernel_group))
+                    choices.append((bench_runner, src_code, self.kernel_group.tile_desc.get_tile_size(), self.kernel_group.tile_desc.vmap.vlane_stride))
                     prevent_infinite_loop += 1
         self.kernel_group.tile_desc.prev_tail_threshold = prev_tail_threshold
         return choices
 
-    def autotune(self, nodes, kernel_name):
+    def autotune(self, *args):
         def get_cycle(choice):
-            bench_runner, src_code, kernel_group = choice
+            bench_runner = choice[0]
             for n_try in range(extension_config.CONFIG_MAX_AUTOTUNE_TRY): # TODO: make simple
                 try:
-                    # bench_runner = self.run_bench(nodes, kernel_name, src_code)
-                    out = bench_runner(validate=extension_config.CONFIG_TORCHSIM_VALIDATION_MODE, autotune=True)
+                    out = bench_runner()
                     return out[-1]
                 except (extension_codecache.SpadOverflowError, RuntimeError) as e:
                     return float("inf")
             return float("inf") # Exceeded maximum number of autotuning attempts
-        choices = self.make_choices(nodes, kernel_name)
+        choices = self.make_choices(*args)
 
         if len(choices) == 0: # can't autotune
             return None
@@ -1628,21 +1633,25 @@ def get_cycle(choice):
         max_idx = results.index(min(results))
         if min(results) == float("inf"):
             raise RuntimeError("Failed to find optimal tile size...")
-        print(f"[Auto-tune] Optimal tile size: {choices[max_idx][2].tile_desc.get_tile_size()}, vlane_stride: {choices[max_idx][2].tile_desc.vmap.vlane_stride}, cycles: {results[max_idx]}")
+        self._log_autotune_result(choices[max_idx], results[max_idx])
         optimal_src_code = choices[max_idx][1]
         return optimal_src_code
 
+    def _log_autotune_result(self, best_choice, best_cycle):
+        print(
+            f"[Auto-tune] Optimal tile size: {list(best_choice[2])}, "
+            f"vlane_stride: {best_choice[3]}, "
+            f"cycles: {best_cycle}"
+        )
+
     def codegen_nodes(self, nodes, kernel_name):
         src_code = super().codegen_nodes(nodes, kernel_name)
         self._prepare_simulator_headers(src_code)
-        if not extension_config.CONFIG_AUTOTUNE or extension_config.CONFIG_BACKENDSIM_SPIKE_ONLY:
-            return src_code
-        else:
+        if extension_config.CONFIG_AUTOTUNE and not extension_config.CONFIG_BACKENDSIM_SPIKE_ONLY:
             optimal_src_code = self.autotune(nodes, kernel_name)
-            if optimal_src_code:
+            if optimal_src_code is not None:
                 return optimal_src_code
-            else:
-                return src_code
+        return src_code
 
     def _prepare_simulator_headers(self, src_code):
         write_path = extension_codecache.get_write_path(src_code)

PyTorchSimFrontend/mlir/mlir_common.py

@@ -408,6 +408,7 @@ def select_vlane_axis(self):
         self.vmap.vlane_split_axis = best_vlane_split_axis
 
     def pad_vlane_tile(self):
+        # FIXME. this doesn't follow tile constraints...
         vlane_split_axis, vlane_stride, vector_lane = self.vmap.vlane_split_axis, self.vmap.vlane_stride, self.vmap.vector_lane
         used_vlane = min(math.ceil(self._tile_size[vlane_split_axis] / vlane_stride), vector_lane)
         padded_size = used_vlane * vlane_stride
@@ -790,7 +791,9 @@ def run_bench(self, nodes, kernel_name, src_code):
                 "vector_lane" : self.vector_lane,
                 "spad_info": self.spad_info,
                 "vlen" : self.vlen,
-                "arg_attributes" : arg_attributes
+                "arg_attributes" : arg_attributes,
+                "validate" : extension_config.CONFIG_TORCHSIM_VALIDATION_MODE,
+                "autotune" : True,
             },
             source_code=src_code,
         )