Add API support for emitting raw tensor slice, update, and reshape. (#59)

stellaraccident · web-flow · commit ce2f5b0d8586 · 2023-09-20T11:20:19.000-07:00
diff --git a/python/shark_turbine/aot/support/ir_utils.py b/python/shark_turbine/aot/support/ir_utils.py
@@ -264,9 +264,19 @@ def build_index_attribute(value: int) -> IntegerAttr:
     return IntegerAttr.get(IndexType.get(), value)
 
 
-def build_index_value(value: int) -> Value:
-    return arith_d.ConstantOp(IndexType.get(), value).result
-
-
-def build_tensor_dim_value(t: Value, dim: int) -> Value:
-    return tensor_d.DimOp(t, build_index_value(dim)).result
+def build_index_value(
+    value: int, constant_cache: Optional[Dict[int, Value]] = None
+) -> Value:
+    if constant_cache is not None and value in constant_cache:
+        return constant_cache[value]
+    index_value = arith_d.ConstantOp(IndexType.get(), value).result
+    if constant_cache is not None:
+        constant_cache[value] = index_value
+    return index_value
+
+
+def build_tensor_dim_value(
+    t: Value, dim: int, constant_cache: Optional[Dict[int, Value]] = None
+) -> Value:
+    dim_value = build_index_value(dim, constant_cache=constant_cache)
+    return tensor_d.DimOp(t, dim_value).result
diff --git a/python/shark_turbine/aot/support/procedural/iree_emitter.py b/python/shark_turbine/aot/support/procedural/iree_emitter.py
@@ -6,7 +6,7 @@
 
 """Python API for IREE's high-level tensor dialects."""
 
-from typing import Any, List, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import functools
 
@@ -36,7 +36,11 @@
 BuildableScalarValue = Union[IrValueScalar, Value]
 BuildableTensorDimDecl = Union[int, Value]
 BuildableTensorType = IrValueTensor
-BuildableIndexType = Union[Value, int]
+BuildableIndexType = Union[BuildableScalarValue, int]
+BuildableIndexLengthType = Union[
+    BuildableTensorDimDecl, Tuple[BuildableTensorDimDecl, BuildableTensorDimDecl]
+]
+BuildableSliceType = Sequence[BuildableIndexLengthType]
 StaticIndexType = int
 
 
@@ -52,9 +56,12 @@ def cast_tensor_value(x: BuildableTensorType) -> IrValueTensor:
     return x
 
 
-def cast_index_value(x: BuildableIndexType) -> Value:
+def cast_index_value(
+    x: BuildableIndexType, *, constant_cache: Optional[Dict[int, Value]] = None
+) -> Value:
+    x = unwrap_intrinsic_value(x)
     if isinstance(x, int):
-        return build_index_value(x)
+        return build_index_value(x, constant_cache=constant_cache)
     else:
         return x
 
@@ -144,6 +151,137 @@ def tensor_empty(
         result.set_dynamic_dim_values(dyn_dim_values)
         return result
 
+    @emitter
+    def tensor_reshape(
+        self, source: BuildableTensorType, *result_dims: BuildableTensorDimDecl
+    ) -> "IrValueTensor":
+        constant_cache: Dict[int, Value] = {}
+        source = cast_tensor_value(source)
+        result_dim_decls, result_dynamic_dims = cast_tensor_dim_decl(result_dims)
+        result_type = RankedTensorType.get(
+            result_dim_decls, source.ir_type.element_type
+        )
+        result_value = flow_d.TensorReshapeOp(
+            result_type,
+            source.ir_value,
+            source.get_only_dynamic_dim_values(constant_cache=constant_cache),
+            result_dynamic_dims,
+        ).result
+        result = IrValueTensor(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":
+        """Extracts a slice of a tensor.
+
+        The given indices must match the rank of the source and each index is
+        interpreted as `(start_index[, length])`, where the `length` is taken
+        to be 1 if only a single value is given for an index.
+        """
+        source = cast_tensor_value(source)
+        source_value = source.ir_value
+        rank = source.rank
+        if len(indices) != rank:
+            raise ValueError(
+                f"Slice indices must match the source rank. Got {len(indices)}, expected {rank}"
+            )
+        # Unpack start_indices and lengths.
+        start_indices: List[BuildableIndexType] = []
+        lengths: List[BuildableIndexType] = []
+        for index_pack in indices:
+            if isinstance(index_pack, (tuple, list)):
+                if len(index_pack) == 2:
+                    start_indices.append(index_pack[0])
+                    lengths.append(index_pack[1])
+                    continue
+            else:
+                start_indices.append(index_pack)
+                lengths.append(1)
+                continue
+            raise ValueError(
+                f"Slice indices expected to be a single value or a 2-tuple. Got {index_pack}"
+            )
+
+        # Process the lengths into a result shape and input length.
+        index_value_cache: Dict[int, Value] = {}
+        length_values: List[Value] = []
+        result_shape: List[int] = []
+        result_dynamic_dims: List[Value] = []
+        for raw_length in lengths:
+            if isinstance(raw_length, int):
+                # Static.
+                result_shape.append(raw_length)
+                if raw_length in index_value_cache:
+                    # Cached.
+                    length_values.append(index_value_cache[raw_length])
+                else:
+                    # Not cached.
+                    length_value = cast_index_value(raw_length)
+                    index_value_cache[raw_length] = length_value
+                    length_values.append(length_value)
+            else:
+                # Dynamic.
+                result_shape.append(ShapedTypeDynamicSizeSentinel)
+                length_value = cast_index_value(raw_length)
+                length_values.append(length_value)
+                result_dynamic_dims.append(length_value)
+        assert len(length_values) == rank
+        assert result_shape.count(ShapedTypeDynamicSizeSentinel) == len(
+            result_dynamic_dims
+        )
+
+        # Process start indices.
+        start_index_values = [cast_index_value(idx) for idx in start_indices]
+        # Emit.
+        result_type = RankedTensorType.get(result_shape, source.ir_type.element_type)
+        constant_cache: Dict[int, Value] = {}
+        result_value = flow_d.TensorSliceOp(
+            result_type,
+            source_value,
+            source.get_only_dynamic_dim_values(constant_cache=constant_cache),
+            start_index_values,
+            length_values,
+            result_dynamic_dims,
+        ).result
+        result = IrValueTensor(result_value, dtype=source.dtype)
+        result.set_dynamic_dim_values(result_dynamic_dims)
+        return result
+
+    @emitter
+    def tensor_update(
+        self,
+        target: BuildableTensorType,
+        update: BuildableTensorType,
+        *start_indices: BuildableIndexType,
+    ) -> "IrValueTensor":
+        """Applies an update to a target at start_indices and returns the mutated target."""
+        constant_cache: Dict[int, Value] = {}
+        target = cast_tensor_value(target)
+        target_dynamic_dims = target.get_only_dynamic_dim_values(
+            constant_cache=constant_cache
+        )
+        update = cast_tensor_value(update)
+        update_dynamic_dims = update.get_only_dynamic_dim_values(
+            constant_cache=constant_cache
+        )
+        start_index_dim_values = [
+            cast_index_value(idx, constant_cache=constant_cache)
+            for idx in start_indices
+        ]
+        result_value = flow_d.TensorUpdateOp(
+            target.ir_value,
+            target_dynamic_dims,
+            start_index_dim_values,
+            update.ir_value,
+            update_dynamic_dims,
+        ).result
+        result = IrValueTensor(result_value, target.dtype)
+        result.set_dynamic_dim_values(target_dynamic_dims)
+        return result
+
     @emitter
     def tensor_splat(
         self,
diff --git a/python/shark_turbine/aot/support/procedural/primitives.py b/python/shark_turbine/aot/support/procedural/primitives.py
@@ -9,9 +9,11 @@
 # operate on instances of these.
 
 from typing import (
+    Dict,
     List,
     Optional,
     Sequence,
+    Tuple,
     Union,
 )
 
@@ -161,7 +163,12 @@ def __repr__(self):
     def resolve_ir_values(self, proc_trace: IrTrace) -> Sequence[Value]:
         return (self.ir_value,)
 
-    def get_dim_value(self, index: int) -> Value:
+    def get_dim_value(
+        self,
+        index: int,
+        *,
+        constant_cache: Optional[Dict[int, Value]] = None,
+    ) -> Value:
         """Gets a dimension as an Index value.
 
         Requires that an InsertionPoint and Location are on the context stack.
@@ -178,9 +185,21 @@ def get_dim_value(self, index: int) -> Value:
             # TODO: Add MLIR API support for creating an insertion point after
             # an operation and use that to set the InsertionPoint to the
             # earliest point.
-            dim_value = build_tensor_dim_value(self.ir_value, index)
+            dim_value = build_tensor_dim_value(
+                self.ir_value, index, constant_cache=constant_cache
+            )
             self._cached_dim_values[index] = dim_value
             return dim_value
         else:
             # Dynamic dim is known.
             return dynamic_dim
+
+    def get_only_dynamic_dim_values(
+        self, *, constant_cache: Optional[Dict[int, Value]] = None
+    ) -> List[Value]:
+        """Returns a list of *only* the dynamic dim Values."""
+        values: List[Value] = []
+        for i, sentinel in enumerate(self._dynamic_dims):
+            if sentinel is not Empty:
+                values.append(self.get_dim_value(i, constant_cache=constant_cache))
+        return values
diff --git a/tests/aot/iree_procedural_test.py b/tests/aot/iree_procedural_test.py
@@ -82,10 +82,110 @@ def foobar(self, x=AbstractIndex, y=AbstractF32):
         inst = BasicModule(context=Context())
         module_str = str(CompiledModule.get_mlir_module(inst))
         print(module_str)
-        self.assertIn("util.global private mutable @_x.global {noinline} : tensor<?x34xf32>", module_str)
+        self.assertIn(
+            "util.global private mutable @_x.global {noinline} : tensor<?x34xf32>",
+            module_str,
+        )
         self.assertIn("%0 = flow.tensor.splat", module_str)
         self.assertIn("util.global.store %0, @_x.global : tensor<?x34xf32>", module_str)
 
+    def testTensorSliceStatic(self):
+        class BasicModule(CompiledModule):
+            def foobar(self, x=AbstractTensor(3, 4)):
+                return IREE.tensor_slice(x, 0, (1, 3))
+
+        inst = BasicModule(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        print(module_str)
+        self.assertIn(
+            "flow.tensor.slice %arg0[%c0, %c1_0 for %c1, %c3] : tensor<3x4xf32> -> tensor<1x3xf32>",
+            module_str,
+        )
+
+    def testTensorSliceDynamicIndex(self):
+        class SliceDynamicIndex(CompiledModule):
+            def foobar(self, x=AbstractIndex):
+                empty = IREE.tensor_empty(x, 16)
+                return IREE.tensor_slice(empty, x, 4)
+
+        inst = SliceDynamicIndex(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        print(module_str)
+        self.assertIn(
+            "flow.tensor.slice %0[%arg0, %c4 for %c1, %c1] : tensor<?x16xf32>{%arg0} -> tensor<1x1xf32>",
+            module_str,
+        )
+
+    def testTensorSliceDynamicLength(self):
+        class SliceDynamicIndex(CompiledModule):
+            def foobar(self, x=AbstractIndex, y=AbstractIndex):
+                empty = IREE.tensor_empty(x, 16)
+                return IREE.tensor_slice(empty, (x, y), 4)
+
+        inst = SliceDynamicIndex(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        print(module_str)
+        self.assertIn(
+            "flow.tensor.slice %0[%arg0, %c4 for %arg1, %c1] : tensor<?x16xf32>{%arg0} -> tensor<?x1xf32>{%arg1}",
+            module_str,
+        )
+
+    def testTensorUpdateStatic(self):
+        class UpdateStatic(CompiledModule):
+            def foobar(
+                self,
+                target=AbstractTensor(4, 4),
+                update=AbstractTensor(2, 2),
+                i=AbstractIndex,
+                j=AbstractIndex,
+            ):
+                return IREE.tensor_update(target, update, i, j)
+
+        inst = UpdateStatic(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        print(module_str)
+        self.assertIn(
+            "flow.tensor.update %arg1, %arg0[%arg2, %arg3] : tensor<2x2xf32> -> %arg0 as tensor<4x4xf32>",
+            module_str,
+        )
+
+    def testTensorUpdateDynamic(self):
+        class UpdateDynamic(CompiledModule):
+            def foobar(
+                self,
+                x=AbstractIndex,
+                y=AbstractIndex,
+                i=AbstractIndex,
+                j=AbstractIndex,
+                value=AbstractF32,
+            ):
+                target = IREE.tensor_empty(x, y)
+                update = IREE.tensor_splat(i, j, value=value, dtype=torch.float32)
+                return IREE.tensor_update(target, update, 2, 2)
+
+        inst = UpdateDynamic(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        print(module_str)
+        self.assertIn(
+            "flow.tensor.update %1, %0[%c2, %c2] : tensor<?x?xf32>{%arg2, %arg3} -> %0 as tensor<?x?xf32>{%arg0, %arg1}",
+            module_str,
+        )
+
+    def testTensorReshape(self):
+        class ReshapeModule(CompiledModule):
+            def foobar(self, x=AbstractIndex, y=AbstractIndex):
+                empty = IREE.tensor_empty(x, 16)
+                reshaped = IREE.tensor_reshape(empty, 1, y, y)
+                return reshaped
+
+        inst = ReshapeModule(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        print(module_str)
+        self.assertIn(
+            "flow.tensor.reshape %0 : tensor<?x16xf32>{%arg0} -> tensor<1x?x?xf32>{%arg1, %arg1}",
+            module_str,
+        )
+
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
