pytorch - 💡(How to fix) Fix [TMA] `TMACompatibilityChecker` fails to detect misalignment when the input pointer has a `constant_offset` [1 pull requests]

import torch, torch._inductor.config as cfg
cfg.triton.use_tensor_descriptor = True
cfg.assume_aligned_inputs = True

def f(x):
    return x[1:] + 1

x = torch.randn(1025, device='cuda')   
torch.compile(f)(x); torch.cuda.synchronize()

---

Traceback (most recent call last):
  File "/opt/pytorch/pytorch/tma_offset.py", line 9, in <module>
    torch.compile(f)(x); torch.cuda.synchronize()
    ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_dynamo/eval_frame.py", line 1141, in compile_wrapper
    result = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/tma_offset.py", line 5, in f
    def f(x):
  File "/opt/pytorch/pytorch/torch/_dynamo/eval_frame.py", line 1415, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_functorch/aot_autograd.py", line 1277, in forward
    return compiled_fn(full_args)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1126, in runtime_wrapper
    result = _codegen_runtime_wrapper(
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_functorch/_aot_autograd/subclass_codegen.py:codegen(runtime_wrapper_orchestration)", line 8, in _runtime_wrapper
  File "/opt/pytorch/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 2551, in __call__
    return self.compiled_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/output_code.py", line 753, in __call__
    return self.current_callable(inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/torchinductor_root/aj/cajf2akezsbggaqpu65jfrz47cqlywlvz2od567mjxlp57f4h6zy.py", line 101, in call
    triton_poi_fused_add_slice_0.run(arg0_1, buf0, 1024, stream=raw_stream0)
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/triton_heuristics.py", line 1990, in run
    self.autotune_to_one_config(*args, **kwargs)
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/triton_heuristics.py", line 1427, in autotune_to_one_config
    timings = self.benchmark_all_configs(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/triton_heuristics.py", line 1391, in benchmark_all_configs
    launcher: self.bench(launcher, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/triton_heuristics.py", line 1234, in bench
    result = benchmarker.benchmark(
             ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/benchmarking.py", line 132, in wrapper
    return fn(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/benchmarking.py", line 249, in benchmark
    return benchmark_fn(self, _callable, warmup=warmup, rep=rep, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/benchmarking.py", line 328, in _default_cuda_bench
    return self.benchmark_gpu(f, warmup=warmup, rep=rep, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/benchmarking.py", line 132, in wrapper
    return fn(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/pytorch/torch/_inductor/runtime/benchmarking.py", line 484, in benchmark_gpu
    torch.cuda.synchronize()
  File "/opt/pytorch/pytorch/torch/cuda/__init__.py", line 1215, in synchronize
    return torch._C._cuda_synchronize()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: misaligned address
Search for `cudaErrorMisalignedAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1