Make tensor updates to a global work. (#64)

dan-garvey · stellaraccident · web-flow · commit 36ec590e0f47 · 2023-09-21T17:49:43.000-07:00
The global and immediate Ir{Scalar|Tensor} hierarchy was unfortunately
divergent. This unifies it and makes the test pass.

---------

Co-authored-by: Stella Laurenzo &lt;stellaraccident@gmail.com&gt;
diff --git a/python/shark_turbine/aot/builtins/jittable.py b/python/shark_turbine/aot/builtins/jittable.py
@@ -41,8 +41,9 @@
 
 from ..support.procedural import (
     CallableIntrinsic,
+    IrImmediateTensor,
+    IrTensor,
     IrTrace,
-    IrValueTensor,
     MaterializedGlobal,
 )
 
@@ -98,7 +99,7 @@ def resolver(py_value: Any, gni: GraphNodeImporter) -> Optional[Value]:
         # Emit a global load and conversion.
         vtensor_type = gni._cc.tensor_to_vtensor_type(py_value)
         loaded_value = util_d.GlobalLoadOp(
-            materialized_global.global_type, materialized_global.symbol_name
+            materialized_global.ir_type, materialized_global.symbol_name
         ).result
         converted_value = Operation.create(
             "torch_c.from_builtin_tensor",
@@ -249,7 +250,7 @@ def resolve_call(self, proc_trace: IrTrace, *py_args, **py_kwargs):
         flat_py_results = []
         for ir_result, pytorch_meta in zip(flat_ir_results, pytorch_meta_results):
             if isinstance(pytorch_meta, TensorMetadata):
-                flat_py_results.append(IrValueTensor(ir_result, pytorch_meta.dtype))
+                flat_py_results.append(IrImmediateTensor(ir_result, pytorch_meta.dtype))
             else:
                 raise TypeError(
                     f"Unknown PyTorch->IREE value mapping for jittable result: {pytorch_meta}->{ir_result}"
@@ -264,7 +265,7 @@ def resolve_call(self, proc_trace: IrTrace, *py_args, **py_kwargs):
             return flat_py_results
 
     def _split_py_arg(self, arg) -> Tuple[Value, Any]:
-        if isinstance(arg, IrValueTensor):
+        if isinstance(arg, IrTensor):
             return arg.ir_value, arg._to_meta_tensor()
 
         raise TypeError(f"Unsupported argument to jittable: {arg}")
diff --git a/python/shark_turbine/aot/support/procedural/base.py b/python/shark_turbine/aot/support/procedural/base.py
@@ -112,6 +112,17 @@ def resolve_assignment(self, proc_trace: "IrTrace", ir_values: Sequence[Value]):
             f"Cannot use {self} as the target of an assignment in a procedural function"
         )
 
+    # Helpers for accessing the ir_value within the current trace.
+    @property
+    def ir_values(self) -> Sequence[Value]:
+        return self.resolve_ir_values(current_ir_trace())
+
+    @property
+    def ir_value(self) -> Value:
+        values = self.ir_values
+        assert len(values) == 1, "Expected arity one intrinsic"
+        return values[0]
+
 
 class CallableIntrinsic(Intrinsic):
     """Intrinsic subclass that supports calls.
@@ -182,7 +193,7 @@ def __repr__(self):
         return f"AbstractTensor({', '.join(str(s) for s in self.size)}, dtype={self.dtype})"
 
     def create_intrinsic(self, ir_value: Value) -> Intrinsic:
-        return IrValueTensor(ir_value, self.dtype)
+        return IrImmediateTensor(ir_value, self.dtype)
 
     def get_ir_type(self, builder: ModuleBuilder) -> IrType:
         element_type = builder.torch_dtype_to_iree_type(self.dtype)
@@ -213,7 +224,7 @@ def __repr__(self):
         return f"AbstractScalar({self.label})"
 
     def create_intrinsic(self, ir_value: Value) -> Intrinsic:
-        return IrValueScalar(ir_value)
+        return IrImmediateScalar(ir_value)
 
     def get_ir_type(self, builder: ModuleBuilder) -> IrType:
         with builder.context:
@@ -249,6 +260,6 @@ def abstractify(tree):
 
 # Circular iports.
 from .primitives import (
-    IrValueScalar,
-    IrValueTensor,
+    IrImmediateScalar,
+    IrImmediateTensor,
 )
diff --git a/python/shark_turbine/aot/support/procedural/globals.py b/python/shark_turbine/aot/support/procedural/globals.py
@@ -35,12 +35,18 @@
 )
 
 from .base import (
-    AbstractTypedef,
+    AbstractScalar,
+    AbstractTensor,
     Intrinsic,
     IrTrace,
     current_ir_trace,
 )
 
+from .primitives import (
+    IrScalar,
+    IrTensor,
+)
+
 ###############################################################################
 # Globals
 ###############################################################################
@@ -121,17 +127,18 @@ def track(self, module_builder: ModuleBuilder, export_namespace: str) -> Any:
                     initialize=self._initialize,
                     mutable=self._mutable,
                 )
-                mapping.value = MaterializedGlobal(
+                mapping.value = IrGlobalTensor(
                     fq_name,
                     self,
                     symbol_name=actual_symbol_name,
                     global_op=global_op,
                     global_type=global_type,
+                    dtype=value.dtype,
                 )
                 logger.debug("TRACK NEW TENSOR(%s): %r", fq_name, mapping)
                 flat_globals.append(mapping.value)
                 continue
-            elif isinstance(value, AbstractTypedef):
+            elif isinstance(value, AbstractTensor):
                 global_type = value.get_ir_type(module_builder)
                 (actual_symbol_name, global_op,) = module_builder.create_typed_global(
                     f"_{fq_name}",
@@ -140,7 +147,26 @@ def track(self, module_builder: ModuleBuilder, export_namespace: str) -> Any:
                     mutable=self._mutable,
                 )
                 flat_globals.append(
-                    MaterializedGlobal(
+                    IrGlobalTensor(
+                        fq_name,
+                        self,
+                        symbol_name=actual_symbol_name,
+                        global_op=global_op,
+                        global_type=global_type,
+                        dtype=value.dtype,
+                    )
+                )
+                continue
+            elif isinstance(value, AbstractScalar):
+                global_type = value.get_ir_type(module_builder)
+                (actual_symbol_name, global_op,) = module_builder.create_typed_global(
+                    f"_{fq_name}",
+                    global_type,
+                    initialize=self._initialize,
+                    mutable=self._mutable,
+                )
+                flat_globals.append(
+                    IrGlobalScalar(
                         fq_name,
                         self,
                         symbol_name=actual_symbol_name,
@@ -158,8 +184,14 @@ def track(self, module_builder: ModuleBuilder, export_namespace: str) -> Any:
             return LiveGlobalCollectionProxy(tree_globals)
 
 
-class MaterializedGlobal(Intrinsic):
-    """Associates a (possibly) materialized global with a name hint and info for the aggregate it is part of."""
+class MaterializedGlobal:
+    """Tags an Ir* that is duck-typed as a global."""
+
+    ...
+
+
+class IrGlobalScalar(IrScalar, MaterializedGlobal):
+    """An IrScalar that is loaded from a global and associated with its aggregate."""
 
     __slots__ = [
         "global_op",
@@ -178,15 +210,65 @@ def __init__(
         global_op: Operation,
         global_type: IrType,
     ):
+        super().__init__(global_type)
+        self.info = info
+        self.export_name = export_name
+        self.symbol_name = symbol_name
+        self.global_op = global_op
+
+    def resolve_ir_values(self, trace: IrTrace) -> Sequence[Value]:
+        with trace.loc, trace.ip:
+            value = util_d.GlobalLoadOp(self.ir_type, self.symbol_name).result
+        return [value]
+
+    def resolve_assignment(self, proc_trace: "IrTrace", ir_values: Sequence[Value]):
+        if len(ir_values) != 1:
+            raise ValueError(
+                f"Can only assign a single value to a global. Got {len(ir_values)}"
+            )
+        source_ir_type = ir_values[0].type
+        if source_ir_type != self.ir_type:
+            raise TypeError(
+                f"Cannot assign to a global with a different type: {self.ir_type} != {source_ir_type}"
+            )
+        with proc_trace.loc, proc_trace.ip:
+            util_d.GlobalStoreOp(ir_values[0], self.symbol_name)
+
+    def __repr__(self):
+        return (
+            f"<IrGlobalScalar {self.export_name} = {self.symbol_name}:{self.ir_type}>"
+        )
+
+
+class IrGlobalTensor(IrTensor, MaterializedGlobal):
+    """An IrScalar that is loaded from a global and associated with its aggregate."""
+
+    __slots__ = [
+        "global_op",
+        "info",
+        "export_name",
+        "symbol_name",
+    ]
+
+    def __init__(
+        self,
+        export_name: str,
+        info: GlobalsDef,
+        *,
+        symbol_name: str,
+        global_op: Operation,
+        global_type: IrType,
+        dtype: torch.dtype,
+    ):
+        super().__init__(global_type, dtype)
         self.info = info
         self.export_name = export_name
         self.symbol_name = symbol_name
         self.global_op = global_op
-        self.global_type = global_type
 
     def resolve_ir_values(self, trace: IrTrace) -> Sequence[Value]:
         with trace.loc, trace.ip:
-            value = util_d.GlobalLoadOp(self.global_type, self.symbol_name).result
+            value = util_d.GlobalLoadOp(self.ir_type, self.symbol_name).result
         return [value]
 
     def resolve_assignment(self, proc_trace: "IrTrace", ir_values: Sequence[Value]):
@@ -195,12 +277,12 @@ def resolve_assignment(self, proc_trace: "IrTrace", ir_values: Sequence[Value]):
                 f"Can only assign a single value to a global. Got {len(ir_values)}"
             )
         source_ir_type = ir_values[0].type
-        if source_ir_type != self.global_type:
+        if source_ir_type != self.ir_type:
             raise TypeError(
-                f"Cannot assign to a global with a different type: {self.global_type} != {source_ir_type}"
+                f"Cannot assign to a global with a different type: {self.ir_type} != {source_ir_type}"
             )
         with proc_trace.loc, proc_trace.ip:
             util_d.GlobalStoreOp(ir_values[0], self.symbol_name)
 
     def __repr__(self):
-        return f"<MaterializedGlobal {self.export_name} = {self.symbol_name}:{self.global_type}>"
+        return f"<MaterializedGlobal {self.export_name} = {self.symbol_name}:{self.ir_type}>"
diff --git a/python/shark_turbine/aot/support/procedural/iree_emitter.py b/python/shark_turbine/aot/support/procedural/iree_emitter.py
@@ -27,15 +27,20 @@
 
 from .base import (
     Intrinsic,
-    IrValueTensor,
-    IrValueScalar,
     current_ir_trace,
     ShapedTypeDynamicSizeSentinel,
 )
 
-BuildableScalarValue = Union[IrValueScalar, Value]
+from .primitives import (
+    IrScalar,
+    IrImmediateScalar,
+    IrTensor,
+    IrImmediateTensor,
+)
+
+BuildableScalarValue = Union[IrScalar, Value]
 BuildableTensorDimDecl = Union[int, Value]
-BuildableTensorType = IrValueTensor
+BuildableTensorType = IrTensor
 BuildableIndexType = Union[BuildableScalarValue, int]
 BuildableIndexLengthType = Union[
     BuildableTensorDimDecl, Tuple[BuildableTensorDimDecl, BuildableTensorDimDecl]
@@ -51,8 +56,8 @@ def cast_scalar_value(x: BuildableScalarValue) -> Value:
     return x
 
 
-def cast_tensor_value(x: BuildableTensorType) -> IrValueTensor:
-    assert isinstance(x, IrValueTensor), f"Expected a tensor but got {type(x)}"
+def cast_tensor_value(x: BuildableTensorType) -> IrTensor:
+    assert isinstance(x, IrTensor), f"Expected a tensor but got {type(x)}"
     return x
 
 
@@ -126,16 +131,16 @@ def wrapper(*args, **kwargs):
 
 class IREEEmitter:
     @emitter
-    def tensor_dim(self, source: BuildableTensorType, index: int) -> "IrValueScalar":
+    def tensor_dim(self, source: BuildableTensorType, index: int) -> "IrScalar":
         """Gets the dimension size of a tensor at a static position."""
         source = cast_tensor_value(source)
         index = cast_static_bounded_index(index, 0, source.rank - 1)
-        return IrValueScalar(source.get_dim_value(index))
+        return IrImmediateScalar(source.get_dim_value(index))
 
     @emitter
     def tensor_empty(
         self, *dims: BuildableTensorDimDecl, dtype: torch.dtype = torch.float32
-    ) -> IrValueTensor:
+    ) -> IrTensor:
         """Constructs a tensor with uninitialized values.
 
         TODO: Support an IREE/raw element type in addition to the torch dtype.
@@ -147,14 +152,14 @@ def tensor_empty(
             raise ValueError(f"Could not map Torch dtype {dtype} to an IREE type")
         tensor_type = RankedTensorType.get(dim_decls, element_type)
         raw_tensor = flow_d.TensorEmptyOp(tensor_type, dyn_dim_values).result
-        result = IrValueTensor(raw_tensor, dtype=dtype)
+        result = IrImmediateTensor(raw_tensor, dtype=dtype)
         result.set_dynamic_dim_values(dyn_dim_values)
         return result
 
     @emitter
     def tensor_reshape(
         self, source: BuildableTensorType, *result_dims: BuildableTensorDimDecl
-    ) -> "IrValueTensor":
+    ) -> "IrTensor":
         constant_cache: Dict[int, Value] = {}
         source = cast_tensor_value(source)
         result_dim_decls, result_dynamic_dims = cast_tensor_dim_decl(result_dims)
@@ -167,14 +172,14 @@ def tensor_reshape(
             source.get_only_dynamic_dim_values(constant_cache=constant_cache),
             result_dynamic_dims,
         ).result
-        result = IrValueTensor(result_value, dtype=source.dtype)
+        result = IrImmediateTensor(result_value, dtype=source.dtype)
         result.set_dynamic_dim_values(result_dynamic_dims)
         return result
 
     @emitter
     def tensor_slice(
         self, source: BuildableTensorType, *indices: BuildableSliceType
-    ) -> "IrValueTensor":
+    ) -> "IrTensor":
         """Extracts a slice of a tensor.
 
         The given indices must match the rank of the source and each index is
@@ -246,7 +251,7 @@ def tensor_slice(
             length_values,
             result_dynamic_dims,
         ).result
-        result = IrValueTensor(result_value, dtype=source.dtype)
+        result = IrImmediateTensor(result_value, dtype=source.dtype)
         result.set_dynamic_dim_values(result_dynamic_dims)
         return result
 
@@ -256,7 +261,7 @@ def tensor_update(
         target: BuildableTensorType,
         update: BuildableTensorType,
         *start_indices: BuildableIndexType,
-    ) -> "IrValueTensor":
+    ) -> "IrTensor":
         """Applies an update to a target at start_indices and returns the mutated target."""
         constant_cache: Dict[int, Value] = {}
         target = cast_tensor_value(target)
@@ -278,7 +283,7 @@ def tensor_update(
             update.ir_value,
             update_dynamic_dims,
         ).result
-        result = IrValueTensor(result_value, target.dtype)
+        result = IrImmediateTensor(result_value, target.dtype)
         result.set_dynamic_dim_values(target_dynamic_dims)
         return result
 
@@ -288,7 +293,7 @@ def tensor_splat(
         *dims: BuildableTensorDimDecl,
         value: BuildableScalarValue,
         dtype: torch.dtype,
-    ) -> "IrValueTensor":
+    ) -> "IrTensor":
         # TODO: Type infer the dtype if missing.
         dim_decls, dyn_dim_values = cast_tensor_dim_decl(dims)
         try:
@@ -302,7 +307,7 @@ def tensor_splat(
             )
         tensor_type = RankedTensorType.get(dim_decls, element_type)
         raw_tensor = flow_d.TensorSplatOp(tensor_type, value, dyn_dim_values).result
-        result = IrValueTensor(raw_tensor, dtype=dtype)
+        result = IrImmediateTensor(raw_tensor, dtype=dtype)
         result.set_dynamic_dim_values(dyn_dim_values)
         return result
 
@@ -314,6 +319,6 @@ def tensor_trace(self, key: str, *ts: BuildableTensorType):
 
 # Circular imports to resolve typing.
 from .primitives import (
-    IrValueScalar,
-    IrValueTensor,
+    IrScalar,
+    IrTensor,
 )
diff --git a/python/shark_turbine/aot/support/procedural/primitives.py b/python/shark_turbine/aot/support/procedural/primitives.py
diff --git a/tests/aot/globals_test.py b/tests/aot/globals_test.py
