ghstack-source-id: c2b9c89
Pull Request resolved: #135313

)

Pull Request resolved: #135313
Approved by: https://github.com/jansel, https://github.com/desertfire
ghstack dependencies: #135312

This reverts commit d383325.

Reverted #135552 on behalf of https://github.com/izaitsevfb due to blocks revert of #135313, internal failures, see D62511427 ([comment](#135552 (comment)))

Fixes #131337

- add `arg_type` for workspace_arg, the type is consistent with the type in `generate_workspace_allocation()`.
- do not generate example tensors for `workspace`, and use `generate_workspace_allocation()` instead.
- add workspace allocation generation code to `kernel_autotune_calls`. e.g.
```python
    workspace = empty_strided_cuda((1280, ), (1, ), torch.uint8)
    workspace.zero_()
    .....
    triton_spl_fused_add_cumprod_0.run(buf2, arg0_1, arg1_1, workspace, 1, 10000, grid=split_scan_grid(1, 10000), stream=stream0)
    del buf2, arg0_1, arg1_1, workspace
```
-  add `empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda` to the header of triton autotune code.

The generated cpp has lines like below, so we also implement a `zero_()` for ` AtenTensorHandle `.

```cpp
    static constexpr int64_t int_array_0[] = {1280L, };
    static constexpr int64_t int_array_1[] = {1L, };
    AtenTensorHandle workspace_handle;
    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_empty_strided(1, int_array_0, int_array_1, cached_torch_dtype_uint8, cached_torch_device_type_cuda,  0, &workspace_handle));

        RAIIAtenTensorHandle workspace(workspace_handle);
        workspace.zero_();
```

- Fix handle grid_fn  for grid computation. Pass in "RBLOCK" to `split_scan_grid`
-  Fix dynamic shapes:
Without the fix we generate code that looks like this `workspace = empty_strided_cuda((32*((255 + s0) // 256), ), (1, ), torch.uint8)` when doing triton autotune and `s0` is not defined.

The solution approach is to use `V.graph.sizevars.size_hint(nbytes)` to realize the workspace size for triton autotune. Note that we only realize it for triton autotune code, but not for the cpp cuda code.

- We also generate slightly different cpp code depending on if `abi_compatible` is turned on.
```cpp
RAIIAtenTensorHandle workspace(workspace_handle);
AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_zero_(workspace.get()));
```
vs

```cpp
    at::Tensor workspace = at::detail::empty_strided_cuda({8L*(c10::div_floor_integer(static_cast<int64_t>((255L + s0)), static_cast<int64_t>(256L))), }, {1L, }, at::kByte, c10::DeviceType::CUDA);
    workspace.zero_();
```

Test Plan:

```
TORCHINDUCTOR_ABI_COMPATIBLE=1 TORCHINDUCTOR_CPP_WRAPPER=1  python test/inductor/test_torchinductor.py -k GPUTests.test_consecutive_split_cumprod_cuda
python test/inductor/test_cuda_cpp_wrapper.py TestCudaWrapper.test_consecutive_split_cumprod_cuda_cuda_wrapper
python test/inductor/test_cuda_cpp_wrapper.py DynamicShapesCudaWrapperCudaTests.test_consecutive_split_cumprod_cuda_dynamic_shapes_cuda_wrapper
TORCHINDUCTOR_ABI_COMPATIBLE=1 python test/inductor/test_cuda_cpp_wrapper.py TestCudaWrapper.test_consecutive_split_cumprod_cuda_cuda_wrapper
TORCHINDUCTOR_CPP_WRAPPER=1  python test/inductor/test_torchinductor.py -k GPUTests.test_consecutive_split_cumprod_cuda
```

Pull Request resolved: #135552
Approved by: https://github.com/desertfire

…y` (#135313)"

This reverts commit 16b37b3.

Reverted #135313 on behalf of https://github.com/izaitsevfb due to breaks internal builds ([comment](#135313 (comment)))

Summary: Reland pytorch#135313 after fixing internal build issues

Test Plan: CI

Differential Revision: D62658837

…y` (#136046)

Summary: Reland #135313 after fixing internal build issues

Test Plan: CI

Differential Revision: D62658837

Pull Request resolved: #136046
Approved by: https://github.com/chenyang78, https://github.com/etaf, https://github.com/jansel

…rch#135313)

Pull Request resolved: pytorch#135313
Approved by: https://github.com/jansel, https://github.com/desertfire
ghstack dependencies: pytorch#135312

This reverts commit d383325.

Reverted pytorch#135552 on behalf of https://github.com/izaitsevfb due to blocks revert of pytorch#135313, internal failures, see D62511427 ([comment](pytorch#135552 (comment)))

…y` (pytorch#135313)"

This reverts commit 16b37b3.

Reverted pytorch#135313 on behalf of https://github.com/izaitsevfb due to breaks internal builds ([comment](pytorch#135313 (comment)))

…y` (pytorch#136046)

Summary: Reland pytorch#135313 after fixing internal build issues

Test Plan: CI

Differential Revision: D62658837

Pull Request resolved: pytorch#136046
Approved by: https://github.com/chenyang78, https://github.com/etaf, https://github.com/jansel

ghstack-source-id: 46b78a4
Pull Request resolved: pytorch/pytorch#135313