Skip to content

Commit

Permalink
[Scheduler] Add the fused task name to auto-scheduled kernels (#418)
Browse files Browse the repository at this point in the history
This PR makes the generated kernel name contained more information.
Previously, the name might look like `hidet_compute_c`, now it might
look like `hidet_fused_softmax_rearrange_c`.
  • Loading branch information
yaoyaoding committed Jan 18, 2024
1 parent f58cc3f commit 1eece05
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion python/hidet/ir/schedulers/cuda/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,13 @@ def schedule_grid_compute(self, node: GridCompute, tensor_map: Dict[TensorNode,
grid_dim: Expr = (prod(node.shape) + block_dim - 1) // block_dim

if self.task is not None:
name = f'{self.task.name}_compute_{node.name}'
from hidet.graph.ops.fusion.fused_operator import FusedTask

if isinstance(self.task, FusedTask):
fused_name = self.task.attrs['fused_ops'].replace(' ', '_')
name = f'fused_{fused_name}_{node.name}'
else:
name = f'{self.task.name}_{node.name}'
else:
name = f'compute_{node.name}'

Expand Down

0 comments on commit 1eece05

Please sign in to comment.