Revert "[dtensor][MTPG] make sharding prop lru cache not shared among…

… threads (pytorch#134294)"

This reverts commit eb15b1a.

Reverted pytorch#134294 on behalf of https://github.com/jeanschmidt due to seems to have introduced https://github.com/pytorch/pytorch/actions/runs/10537099590/job/29201744658 ([comment](pytorch#134294 (comment)))

test/distributed/_tensor/test_dtensor_ops.py

@@ -353,6 +353,7 @@ def wrapped(fn):
     xfail("nn.functional.pdist"),
     xfail("nn.functional.pixel_shuffle"),
     xfail("nn.functional.pixel_unshuffle"),
+    xfail("nn.functional.poisson_nll_loss"),
     xfail("nn.functional.prelu"),
     xfail("nn.functional.relu6"),
     xfail("nn.functional.rrelu"),

torch/distributed/_tensor/_sharding_prop.py

@@ -1,5 +1,4 @@
 # mypy: allow-untyped-defs
-import threading
 from functools import lru_cache
 from itertools import chain
 from typing import Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
@@ -38,17 +37,6 @@ def _length(obj) -> int:
     return len(obj)
 
 
-class LocalLRUCache(threading.local):
-    def __init__(self, user_function: Callable) -> None:
-        self.cache = lru_cache(None)(user_function)
-
-    def __call__(self, *args, **kwargs) -> object:
-        return self.cache(*args, **kwargs)
-
-    def cache_info(self):
-        return self.cache.cache_info()
-
-
 class ShardingPropagator:
     def __init__(self) -> None:
         self.op_to_rules: Dict[OpOverload, Callable[[OpSchema], OutputSharding]] = {}
@@ -58,9 +46,7 @@ def __init__(self) -> None:
         ] = {}
         # op map to save static argnum to decide to reuse sharding prop cache or re-run sharding prop
         self.op_to_schema_info: Dict[OpOverload, RuntimeSchemaInfo] = {}
-        self.propagate_op_sharding = LocalLRUCache(
-            self.propagate_op_sharding_non_cached
-        )
+        self.propagate_op_sharding = lru_cache(None)(self.propagate_op_sharding_non_cached)  # type: ignore[method-assign]
         # op map to save indices of shape (and stride) args which may need to be modified in sharding prop
         self.op_to_shape_and_stride_idx: Dict[
             OpOverload, Union[int, Tuple[int, int]]
@@ -197,9 +183,7 @@ def propagate(self, op_info: OpInfo) -> None:
         if op_info.schema.has_symints:
             output_sharding = self.propagate_op_sharding_non_cached(op_info.schema)
         else:
-            output_sharding = cast(
-                OutputSharding, self.propagate_op_sharding(op_info.schema)
-            )
+            output_sharding = self.propagate_op_sharding(op_info.schema)
         op_info.output_sharding = output_sharding
 
     def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputSharding: