Add tests.

Author: stellaraccident

python/shark_turbine/ops/iree.py

@@ -43,8 +43,7 @@ def _emit_tensor_trace(kb: KernelBuilder, key: str, ts: list[Value]):
 
 @CustomOp.register(library=IREE_LIBRARY)
 class trace_tensor(CustomOp):
-    name = "trace_tensor"
-    signature = "(str trace_key, Tensor tensor) -> ()"
+    signature = "trace_tensor(str trace_key, Tensor tensor) -> ()"
 
     def select(self, ksel: KernelSelection):
         ksel.attr_str(0)
@@ -57,8 +56,7 @@ def generate(self, ksel: KernelSelection, kb: KernelBuilder):
 
 @CustomOp.register(library=IREE_LIBRARY)
 class trace_tensors(CustomOp):
-    name = "trace_tensors"
-    signature = "(str trace_key, Tensor[] tensors) -> ()"
+    signature = "trace_tensors(str trace_key, Tensor[] tensors) -> ()"
 
     def select(self, ksel: KernelSelection):
         ksel.attr_str(0)

python/shark_turbine/runtime/op_reg/base.py

@@ -13,6 +13,7 @@
 from abc import ABC, abstractmethod, abstractproperty
 import functools
 import logging
+import re
 import textwrap
 
 import torch
@@ -132,8 +133,8 @@ def __init__(
         register_meta: bool,
         register_impl: bool,
     ):
-        name = self.name
-        fq_schema = f"{name}{self.signature}"
+        fq_schema = self.signature
+        name = _extract_name_from_signature(fq_schema)
         library.define(fq_schema)
         self.library = library
         self.cache_key_base = f"{library.ns}.{library.kind}::{name}"
@@ -156,17 +157,15 @@ def __init__(
         fq_name = f"{library.ns}.{name}"
         ALL_CUSTOM_OP_REGS[fq_name] = self
 
-    @abstractproperty
-    def name(self) -> str:
-        """Name of the operation."""
-        ...
-
     @abstractproperty
     def signature(self) -> str:
         """PyTorch function signature.
 
-        This excludes the name, which will come from the `name` property
-        and be prepended to make a full PyTorch schema.
+        This is in the normal PyTorch kernel registration form. For example:
+
+        ```
+        my_op(Tensor t) -> Tensor
+        ```
         """
         ...
 
@@ -772,3 +771,13 @@ def handler(*args):
         return eager_dispatch(ksel)
 
     return handler
+
+
+_SIGNATURE_NAME_PATTERN = re.compile(r"^([^(]+)")
+
+
+def _extract_name_from_signature(sig: str) -> str:
+    m = re.match(_SIGNATURE_NAME_PATTERN, sig)
+    if not m:
+        raise ValueError(f"Expected signature of form `name() -> (). Got: {sig}")
+    return m.group(1)

python/shark_turbine/transforms/general/custom_op_expansion.py

@@ -43,8 +43,11 @@
 
 
 class ExpandCustomOpsPass(Pass):
-    def __init__(self, root_op: Operation):
+    def __init__(
+        self, root_op: Operation, reg: dict[str, CustomOp] = ALL_CUSTOM_OP_REGS
+    ):
         super().__init__(root_op)
+        self.reg = reg
         # Track pending deletions in a dict to preserve order and unique.
         self.ops_to_delete: dict[Operation, None] = {}
         self.type_converter = NativeTypeConverter(root_op.context)
@@ -76,7 +79,7 @@ def expand_func(self, func_op: Operation):
                     custom_op_name = StringAttr(op.attributes["name"]).value
                     if custom_op_name.startswith(name_prefix):
                         local_name = custom_op_name[len(name_prefix) :]
-                        custom_op_reg = ALL_CUSTOM_OP_REGS.get(local_name)
+                        custom_op_reg = self.reg.get(local_name)
                         if custom_op_reg is not None:
                             self.expand_custom_op(custom_op_reg, op)
 
@@ -172,7 +175,9 @@ def attr_str(self, arg: int) -> AttrArg:
         return desc
 
     def return_tensor(self, t: Tensor) -> TensorArg:
-        raise NotImplementedError("NYI: return_tensor")
+        desc = TensorArg(t)
+        self.result_descs.append(desc)
+        return desc
 
 
 def _get_constant_str_from_value(v: Value) -> str:
@@ -224,6 +229,7 @@ def __init__(
         )
         self.location = location
         self.torch_op = torch_op
+        self.type_converter = type_converter
 
     def yield_results(self, *results: Value):
         """Yields results of the kernel computation."""
@@ -234,5 +240,9 @@ def yield_results(self, *results: Value):
                 torch_op_results
             ), f"Mismatched yield_results with custom op results"
             for new_result, old_result in zip(results, torch_op_results):
+                torch_type = old_result.type
+                new_result = self.type_converter.materialize_native_to_torch(
+                    new_result, torch_type
+                )
                 old_result.replace_all_uses_with(new_result)
         self.yielded = True

tests/dynamo/type_conversion_test.py

@@ -24,15 +24,19 @@ def testPrimitives(self):
         self._compareNative("!torch.int", "i64")
         self._compareNative("!torch.float", "f64")
 
+    def testSigned(self):
+        self._compareNative("!torch.bool", "i1", signless=False)
+        self._compareNative("!torch.int", "si64", signless=False)
+
     def testValueTensors(self):
         self._compareNative("!torch.vtensor<[2, 2],f32>", "tensor<2x2xf32>")
         self._compareNative("!torch.vtensor<[?, ?],f32>", "tensor<?x?xf32>")
         self._compareNative("!torch.vtensor<[],f32>", "tensor<f32>")
 
-    def _compareNative(self, torch_str: str, native_str: str):
+    def _compareNative(self, torch_str: str, native_str: str, *, signless: bool = True):
         with self.conv._context:
             torch_type = IrType.parse(torch_str)
-        native_type = self.conv.torch_type_to_native(torch_type)
+        native_type = self.conv.torch_type_to_native(torch_type, signless=signless)
         self.assertEqual(str(native_type), native_str)
 
 

tests/runtime/op_reg/kernel_aot_test.py

@@ -43,14 +43,15 @@ class KernelRegTest(unittest.TestCase):
     def testTrace(self):
         mlp = MLP()
         prog = aot.export(mlp, torch.empty(97, 8, dtype=torch.float32))
-        # print("ORIGINAL EXPORTED:")
-        # print(prog.print_readable())
 
         p = ExpandCustomOpsPass(prog.mlir_module)
         p.run()
 
         print("CUSTOM OP CONVERTED:")
-        print(prog.mlir_module)
+        module_asm = str(prog.mlir_module)
+        self.assertIn('flow.tensor.trace "LAYER0"', module_asm)
+        self.assertIn('flow.tensor.trace "LAYER1"', module_asm)
+        self.assertIn('flow.tensor.trace "LAYER3"', module_asm)
 
     def testEager(self):
         mlp = MLP()

tests/runtime/op_reg/kernel_reg_test.py

@@ -15,11 +15,10 @@
 
 
 class KernelRegTest(unittest.TestCase):
-    def testSimple(self):
+    def testRegistrationDispatchAndCache(self):
         @CustomOp.register
         class identity(CustomOp):
-            name = "test_identity"
-            signature = "(Tensor self) -> Tensor"
+            signature = "test_identity(Tensor self) -> Tensor"
 
             def select(self, ksel: KernelSelection):
                 x = ksel.arg_tensor(0)

tests/transforms/general/custom_op_expansion_test.py

@@ -0,0 +1,112 @@
+# Copyright 2024 Advanced Micro Devices, Inc
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import logging
+from pathlib import Path
+import torch
+import unittest
+
+from shark_turbine.transforms.general.custom_op_expansion import ExpandCustomOpsPass
+from shark_turbine.runtime.op_reg import (
+    def_library,
+    CustomOp,
+    KernelBuilder,
+    KernelSelection,
+)
+
+from shark_turbine.support.ir_imports import (
+    Context,
+    Module,
+)
+
+
+class PassTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.lib = def_library("expand_custom_op_pass_test")
+        CustomOp.register(library=cls.lib)(IdentityOp)
+        CustomOp.register(library=cls.lib)(PrintStringAttrOp)
+        CustomOp.register(library=cls.lib)(IntArgOp)
+
+    def testTensorArgReturn(self):
+        m = self.run_test_case("custom_op_simple.mlir")
+        m_asm = str(m)
+        self.assertNotIn("torch.operator", m_asm)
+        self.assertIn(
+            "%0 = torch_c.to_builtin_tensor %arg0 : !torch.vtensor<[97,8],f32> -> tensor<97x8xf32>",
+            m_asm,
+        )
+        self.assertIn(
+            "%1 = torch_c.from_builtin_tensor %0 : tensor<97x8xf32> -> !torch.vtensor<[97,8],f32>",
+            m_asm,
+        )
+        print(m_asm)
+
+    def testStringAttrArg(self):
+        global _TEST_STRING_ATTR
+        _TEST_STRING_ATTR = ""
+        m = self.run_test_case("custom_op_string_attr.mlir")
+        m_asm = str(m)
+        self.assertEqual(_TEST_STRING_ATTR, "TEST_VALUE")
+        self.assertNotIn("torch.operator", m_asm)
+        print(m_asm)
+
+    def testIntArg(self):
+        global _TEST_STRING_ATTR
+        _TEST_STRING_ATTR = ""
+        with self.assertRaisesRegex(NotImplementedError, "arg_int"):
+            self.run_test_case("custom_op_int_arg.mlir")
+
+    def run_test_case(self, file_name: str):
+        p = Path(__file__).resolve().parent / "testdata" / file_name
+        contents = p.read_text()
+        with Context() as ctx:
+            m = Module.parse(contents)
+        p = ExpandCustomOpsPass(m.operation)
+        p.run()
+        print(f"TEST CASE {file_name}:\n{m}")
+        m.operation.verify()
+        return m
+
+
+class IdentityOp(CustomOp):
+    signature = "identity_tensor(Tensor t) -> Tensor"
+
+    def select(self, ksel: KernelSelection):
+        x = ksel.arg_tensor(0)
+        ksel.return_tensor(x.t)
+
+    def generate(self, ksel: KernelSelection, kb: KernelBuilder):
+        kb.yield_results(kb.arg_bindings[0])
+
+
+class PrintStringAttrOp(CustomOp):
+    signature = "print_string_attr(str key) -> ()"
+
+    def select(self, ksel: KernelSelection):
+        ksel.attr_str(0)
+
+    def generate(self, ksel: KernelSelection, kb: KernelBuilder):
+        global _TEST_STRING_ATTR
+        _TEST_STRING_ATTR = str(ksel.arg_descs[0].v)
+        print("CAPTURED STRING ATTR:", _TEST_STRING_ATTR)
+        kb.yield_results()
+
+
+class IntArgOp(CustomOp):
+    signature = "int_arg(int t) -> ()"
+
+    def select(self, ksel: KernelSelection):
+        x = ksel.arg_int(0)
+        ksel.return_int()
+
+    def generate(self, ksel: KernelSelection, kb: KernelBuilder):
+        kb.yield_results(kb.arg_bindings[0])
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    unittest.main()

tests/transforms/general/testdata/custom_op_int_arg.mlir

@@ -0,0 +1,9 @@
+builtin.module {
+
+func.func @forward() {
+  %i = torch.constant.int 1000
+  torch.operator "torch.expand_custom_op_pass_test.int_arg"(%i) : (!torch.int) -> ()
+  return
+}
+
+}

tests/transforms/general/testdata/custom_op_simple.mlir

@@ -0,0 +1,8 @@
+builtin.module {
+
+func.func @forward(%arg0: !torch.vtensor<[97,8],f32>) -> !torch.vtensor<[97,8],f32> {
+  %0 = torch.operator "torch.expand_custom_op_pass_test.identity_tensor"(%arg0) : (!torch.vtensor<[97,8],f32>) -> (!torch.vtensor<[97,8],f32>)
+  return %0 : !torch.vtensor<[97,8],f32>
+}
+
+}

tests/transforms/general/testdata/custom_op_string_attr.mlir

@@ -0,0 +1,9 @@
+builtin.module {
+
+func.func @forward() {
+  %str = torch.constant.str "TEST_VALUE"
+  torch.operator "torch.expand_custom_op_pass_test.print_string_attr"(%str) : (!torch.str) -> ()
+  return  
+}
+
+}