Repository: pytorch
Issue: `torch.profiler.record_function` doesn't register kernels from `backward` function
Labels: oncall: profiler

### 🐛 Describe the bug

`__profile_kernel_of_func` (`record_function` label) shows zero timings for XPU (maybe for `CUDA` the situation is the same, but I have no way to check) unless `record_function` is used inside `backward` function.

```python
import torch
from torch.profiler import profile, ProfilerActivity, record_function

cache_size = 256
device = "xpu"

class _attention(torch.autograd.Function):

    @staticmethod
    def forward(ctx, cache):
        ctx.save_for_backward(cache)
        return cache

    @staticmethod
    def backward(ctx, triton_do):
        cache = ctx.saved_tensors
        # with record_function("__profile_kernel_of_func"): <- using this you can get the necessary timings
        cache[0].zero_()
        return cache
    

attention = _attention.apply

cache = torch.randn((128, 128), dtype=torch.float32, device=device, requires_grad=True)
triton_o = attention(cache)

triton_do = torch.randn_like(triton_o)
triton_fn = lambda: triton_o.backward(triton_do, retain_graph=True)

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.XPU]) as prof:
    with record_function("__profile_kernel_of_func"):
        triton_fn()
    torch.xpu.synchronize()

print(prof.events())

```

Output:
```

# case1 - `record_function` is not used in `backward` function and one can see that timings are zero

-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg      Self XPU    Self XPU %     XPU total  XPU time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                               __profile_kernel_of_func        90.70%     193.820ms        90.70%     193.820ms     193.820ms       0.000us         0.00%       0.000us       0.000us             1  
autograd::engine::evaluate_function: _attentionBackw...         0.01%      20.711us         3.38%       7.220ms       7.220ms       0.000us         0.00%       8.640us       8.640us             1  
                                     _attentionBackward         0.04%      86.246us         3.37%       7.199ms       7.199ms       0.000us         0.00%       8.640us       8.640us             1  
                                            aten::zero_         0.03%      53.549us         3.33%       7.113ms       7.113ms       0.000us         0.00%       8.640us       8.640us             1  
                                            aten::fill_         3.17%       6.772ms         3.30%       7.059ms       7.059ms       8.640us        50.00%       8.640us       8.640us             1  
                                  urEnqueueKernelLaunch         0.13%     287.035us         0.13%     287.035us     287.035us       0.000us         0.00%       0.000us       0.000us             1  
at::native::xpu::VectorizedElementwiseKernel<4, at::...         0.00%       0.000us         0.00%       0.000us       0.000us       8.640us        50.00%       8.640us       8.640us             1  
autograd::engine::evaluate_function: torch::autograd...         0.00%       5.280us         5.93%      12.663ms      12.663ms       0.000us         0.00%       8.640us       8.640us             1  
                        torch::autograd::AccumulateGrad         0.02%      34.944us         5.92%      12.658ms      12.658ms       0.000us         0.00%       8.640us       8.640us             1  
                                aten::new_empty_strided         0.00%       7.971us         0.01%      24.983us      24.983us       0.000us         0.00%       0.000us       0.000us             1  
                                    aten::empty_strided         0.01%      17.012us         0.01%      17.012us      17.012us       0.000us         0.00%       0.000us       0.000us             1  
                                            aten::copy_         5.86%      12.514ms         5.90%      12.598ms      12.598ms       8.640us        50.00%       8.640us       8.640us             1  
                                  urEnqueueKernelLaunch         0.04%      83.934us         0.04%      83.934us      83.934us       0.000us         0.00%       0.000us       0.000us             1  
at::native::xpu::VectorizedElementwiseKernel<4, at::...         0.00%       0.000us         0.00%       0.000us       0.000us       8.640us        50.00%       8.640us       8.640us             1  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 213.703ms
Self XPU time total: 17.280us


case #2 - `record_function` is used in `backward` function

-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg      Self XPU    Self XPU %     XPU total  XPU time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                               __profile_kernel_of_func        90.60%     218.883ms        90.60%     218.883ms     218.883ms       0.000us         0.00%       0.000us       0.000us             1  
autograd::engine::evaluate_function: _attentionBackw...         0.01%      20.057us         3.60%       8.690ms       8.690ms       0.000us         0.00%       8.320us       8.320us             1  
                                     _attentionBackward         0.03%      81.833us         3.59%       8.670ms       8.670ms       0.000us         0.00%       8.320us       8.320us             1  
                               __profile_kernel_of_func         0.10%     230.601us         3.55%       8.588ms       8.588ms       0.000us         0.00%       8.320us       8.320us             1  
                                            aten::zero_         0.02%      54.334us         3.46%       8.358ms       8.358ms       0.000us         0.00%       8.320us       8.320us             1  
                                            aten::fill_         3.30%       7.984ms         3.44%       8.304ms       8.304ms       8.320us        49.06%       8.320us       8.320us             1  
                                  urEnqueueKernelLaunch         0.13%     319.102us         0.13%     319.102us     319.102us       0.000us         0.00%       0.000us       0.000us             1  
at::native::xpu::VectorizedElementwiseKernel<4, at::...         0.00%       0.000us         0.00%       0.000us       0.000us       8.320us        49.06%       8.320us       8.320us             1  
autograd::engine::evaluate_function: torch::autograd...         0.00%       6.131us         5.80%      14.019ms      14.019ms       0.000us         0.00%       8.640us       8.640us             1  
                        torch::autograd::AccumulateGrad         0.02%      42.298us         5.80%      14.013ms      14.013ms       0.000us         0.00%       8.640us       8.640us             1  
                                aten::new_empty_strided         0.00%       9.754us         0.01%      26.166us      26.166us       0.000us         0.00%       0.000us       0.000us             1  
                                    aten::empty_strided         0.01%      16.412us         0.01%      16.412us      16.412us       0.000us         0.00%       0.000us       0.000us             1  
                                            aten::copy_         5.73%      13.855ms         5.77%      13.945ms      13.945ms       8.640us        50.94%       8.640us       8.640us             1  
                                  urEnqueueKernelLaunch         0.04%      90.206us         0.04%      90.206us      90.206us       0.000us         0.00%       0.000us       0.000us             1  
at::native::xpu::VectorizedElementwiseKernel<4, at::...         0.00%       0.000us         0.00%       0.000us       0.000us       8.640us        50.94%       8.640us       8.640us             1  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 241.593ms
Self XPU time total: 16.960us
```

### Versions

Pytorch pin: `1e881ceecfe80532206ca4e0acb64391fab8b935`.

cc @robieta @chaekit @guotuofeng @guyang3532 @dzhulgakov @davidberard98 @briancoutinho @sraikund16 @sanrise