[inductor] share cse cache during vectorized indirect load

[zhuhaozhe]

Fix #123502

swap_buffer in not needed in vectorized indirect load, remove it to share cse buffer.

auto tmp8 =
[&]
{
    __at_align__ std::array<int64_t, 16> tmpbuf;
    tmp7.store(tmpbuf.data());
    return tmpbuf;
}
()
;
//
// other codes
//
// also store tmp7 here (redundant tmp16)
auto tmp16 =
[&]
{
    __at_align__ std::array<int64_t, 16> tmpbuf;
    tmp7.store(tmpbuf.data());
    return tmpbuf;
}
()
;

Stack from ghstack (oldest at bottom):

• [inductor] share more cse cache during swap buffer #124921
• -> [inductor] share cse cache during vectorized indirect load #124597

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124597

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit dda79ae with merge base 59a1f1f ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  inductor/test_cudagraph_trees.py::CudaGraphTreeTests::test_mutation_cudagraph_managed_tensors_config_backend_cudagraphs

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jgong5]

The "cse" object has cache, reduction_cache and store_cache. A single cse_cache cannot represent all of them. Maybe it is better to pass cse directly as to override the new "cse" after the buffers are swapped?

[jgong5]

When new variables are defined in the scope, they cannot be used outside the scope afterwards. So, it is not safe to reuse the existing self.cse.cache here. I guess we want to create a wrapper around the cache abstraction that supports scoping, allowing variables and expressions defined within a scope to be valid only within that scope while still leveraging the existing cache from the outer scope before it. How about we create something like below?

def scope_cse(cse):
    new_cse = cse.clone()
    new_cse.cache = ScopedDict(cse.cache)
    new_cse.reduction_cache = ScopedDict(cse.reduction_cache)
    new_cse.store_cache = ScopedDict(cse.store_cache)
    return new_cse

Here the ScopedDict allows read-only access to the original dict while uses a separate dict to hold newly-added items:

class ScopedDict:
    def __init__(self, original_dict):
        self.original_dict = original_dict
        self.new_items = {}
    
    def __getitem__(self, key):
        if key in self.new_items:
            return self.new_items[key]
        return self.original_dict[key]
    
    def __setitem__(self, key, value):
        self.new_items[key] = value
    
    def __contains__(self, key):
        return key in self.new_items or key in self.original_dict
    
    def get(self, key, default=None):
        if key in self.new_items:
            return self.new_items[key]
        return self.original_dict.get(key, default)

[zhuhaozhe]

When new variables are defined in the scope, they cannot be used outside the scope afterwards. So, it is not safe to reuse the existing self.cse.cache here. I guess we want to create a wrapper around the cache abstraction that supports scoping, allowing variables and expressions defined within a scope to be valid only within that scope while still leveraging the existing cache from the outer scope before it. How about we create something like below?

def scope_cse(cse):
    new_cse = cse.clone()
    new_cse.cache = ScopedDict(cse.cache)
    new_cse.reduction_cache = ScopedDict(cse.reduction_cache)
    new_cse.store_cache = ScopedDict(cse.store_cache)
    return new_cse

Here the ScopedDict allows read-only access to the original dict while uses a separate dict to hold newly-added items:

class ScopedDict:
    def __init__(self, original_dict):
        self.original_dict = original_dict
        self.new_items = {}
    
    def __getitem__(self, key):
        if key in self.new_items:
            return self.new_items[key]
        return self.original_dict[key]
    
    def __setitem__(self, key, value):
        self.new_items[key] = value
    
    def __contains__(self, key):
        return key in self.new_items or key in self.original_dict
    
    def get(self, key, default=None):
        if key in self.new_items:
            return self.new_items[key]
        return self.original_dict.get(key, default)

Sounds better, I will try this way.

[zhuhaozhe]

Hi, @jgong5. For this issue #123502, your recommendation for ScopedDict should not work because

            auto tmp8 =
            [&]
            {
                __at_align__ std::array<int64_t, 16> tmpbuf;
                tmp7.store(tmpbuf.data());
                return tmpbuf;
            }
            ()
            ;
            auto tmp9 =
            [&]
            {
                __at_align__ std::array<float, 16> tmpbuf;
                #pragma GCC unroll 16
                for (long x0_inner = 0; x0_inner < 16; x0_inner++)
                {
                    tmpbuf[x0_inner] = out_ptr0[static_cast<long>(tmp8[x0_inner])];
                }
                return at::vec::Vectorized<float>::loadu(tmpbuf.data());
            }
            ()
            ;

the tmp8 is created inside this scope here. But actually I observed the tmp8 should be able to be used outside because the buffer here in csevar = self.cse.generate(buffer, code) is not inside the lambda function (of tmp9) scope.

For my current understanding we may not need to always put a swap_buffer after a [&], for _load_or_store_non_contiguous ,should it be safe to remove swap buffer here?

[zhuhaozhe]

As for the ScopedDict, it indeed help.
I found a case test_AllenaiLongformerBase_repro_cpu and make it as an UT in another PR #124921.
It will only generate tmp_x=static_cast<int>(256) once after ScopedDict impl, and it will be 3 times before it.

[zhuhaozhe]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here