[NFC] use register_value_caster for Index and F16Vector

.github/workflows/build_test_tutorial.yml

@@ -0,0 +1,99 @@
+# Part of the LLVM Project, 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
+# Copyright (c) 2024.
+
+name: "Test tutorial"
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.number || github.sha }}-test_tutorial
+  cancel-in-progress: true
+
+jobs:
+
+  test-tutorial:
+
+    strategy:
+      fail-fast: false
+      matrix:
+        runs-on: [
+          "ubuntu-22.04",
+          "ubuntu-22.04-arm",
+          "macos-14",
+          "windows-2022"
+        ]
+        python-version: [
+          "3.10", "3.11", "3.12",
+          "3.13", "3.14", "3.14t"
+        ]
+        include: [
+          {runs-on: "ubuntu-22.04", name: "ubuntu_x86_64", os: "ubuntu"},
+          {runs-on: "ubuntu-22.04-arm", name: "ubuntu_aarch64", os: "ubuntu"},
+          {runs-on: "windows-2022", name: "windows_amd64", os: "windows"},
+          {runs-on: "macos-14", name: "macos_arm64", os: "macos"},
+        ]
+        exclude:
+          # <frozen importlib._bootstrap>:491: Warning: Numpy built with MINGW-W64 on Windows 64 bits is experimental, and only available for testing. You are advised not to use it for production.
+          - runs-on: windows-2022
+            python-version: "3.14"
+
+          - runs-on: windows-2022
+            python-version: "3.14t"
+
+          - runs-on: macos-14
+            python-version: "3.10"
+
+          - runs-on: macos-14
+            python-version: "3.11"
+
+          - runs-on: macos-14
+            python-version: "3.12"
+
+          - runs-on: macos-14
+            python-version: "3.13"
+
+          - runs-on: macos-14
+            python-version: "3.14"
+
+    runs-on: ${{ matrix.runs-on }}
+
+    name: "Test tutorial ${{ matrix.name }} ${{ matrix.python-version }}"
+
+    defaults:
+      run:
+        shell: bash
+
+    steps:
+      - name: "Check out repository"
+        uses: actions/checkout@v4.2.2
+        with:
+          submodules: false
+
+      - name: "Install Python"
+        uses: actions/setup-python@v6.0.0
+        with:
+          python-version: "${{ matrix.python-version }}"
+
+      - name: "Install requirements"
+        run: python -m pip install -r requirements.txt
+
+      - name: "Build tutorial"
+        run: |
+
+          cmake -B build -S $PWD -DCMAKE_PREFIX_PATH=$(python -m mlir_wheel --root-dir) -DLLVM_EXTERNAL_LIT=$(which lit)
+          cmake --build build --target tutorial-opt
+
+      - name: "Test tutorial"
+        run: |
+
+          cmake --build build --target check-tutorial
+

python/tiny/ch1/__init__.py

@@ -2,8 +2,3 @@
 
 from .dsl import Index, F16Vector, Ptr
 from .dsl import print_ir, compile_and_print
-
-__all__ = [
-    "Index", "F16Vector", "Ptr",
-    "print_ir", "compile_and_print",
-]

python/tiny/ch1/dsl.py

@@ -1,22 +1,22 @@
 """DSL wrapper types for Tiny dialect."""
 
 from mlir.ir import (
-    Value, Type, Operation,
-    IndexType, VectorType, F16Type,
-    IntegerAttr, DenseElementsAttr,
+    Value,
+    Type,
+    Operation,
+    IndexType,
+    VectorType,
+    F16Type,
+    IntegerAttr,
+    DenseElementsAttr,
+    register_value_caster,
 )
 import numpy as np
 
 
-class Index:
+@register_value_caster(IndexType.static_typeid)
+class Index(Value):
     """Wrapper for index-typed SSA values."""
-    _value: Value
-
-    @staticmethod
-    def _wrap(value: Value) -> "Index":
-        idx = Index()
-        idx._value = value
-        return idx
 
     @staticmethod
     def constant(val: int) -> "Index":
@@ -28,31 +28,34 @@ def constant(val: int) -> "Index":
             results=[idx_type],
             attributes={"value": attr},
         )
-        return Index._wrap(op.result)
+        return op.result
 
     def _binop(self, other: "Index", op_name: str) -> "Index":
         op = Operation.create(
             f"tiny.{op_name}",
             results=[IndexType.get()],
-            operands=[self._value, other._value],
+            operands=[self, other],
         )
-        return Index._wrap(op.result)
+        return op.result
 
-    def __add__(self, other): return self._binop(other, "addi")
-    def __sub__(self, other): return self._binop(other, "subi")
-    def __mul__(self, other): return self._binop(other, "muli")
-    def __floordiv__(self, other): return self._binop(other, "divi")
+    def __add__(self, other):
+        return self._binop(other, "addi")
 
+    def __sub__(self, other):
+        return self._binop(other, "subi")
 
-class F16Vector:
-    """Wrapper for vector<Nxf16> SSA values."""
-    _value: Value
+    def __mul__(self, other):
+        return self._binop(other, "muli")
 
-    @staticmethod
-    def _wrap(value: Value) -> "F16Vector":
-        vec = F16Vector()
-        vec._value = value
-        return vec
+    def __floordiv__(self, other):
+        return self._binop(other, "divi")
+
+    def __repr__(self):
+        return repr(self).replace("Value", "Index")
+
+
+class F16Vector(Value):
+    """Wrapper for vector<Nxf16> SSA values."""
 
     @staticmethod
     def constant(vals: list[float], size: int = None) -> "F16Vector":
@@ -66,34 +69,52 @@ def constant(vals: list[float], size: int = None) -> "F16Vector":
             results=[vec_type],
             attributes={"value": attr},
         )
-        return F16Vector._wrap(op.result)
+        return op.result
 
     def _binop(self, other: "F16Vector", op_name: str) -> "F16Vector":
         op = Operation.create(
             f"tiny.{op_name}",
-            results=[self._value.type],
-            operands=[self._value, other._value],
+            results=[self.type],
+            operands=[self, other],
         )
-        return F16Vector._wrap(op.result)
+        return op.result
 
-    def __add__(self, other): return self._binop(other, "addf")
-    def __sub__(self, other): return self._binop(other, "subf")
-    def __mul__(self, other): return self._binop(other, "mulf")
-    def __truediv__(self, other): return self._binop(other, "divf")
+    def __add__(self, other):
+        return self._binop(other, "addf")
+
+    def __sub__(self, other):
+        return self._binop(other, "subf")
+
+    def __mul__(self, other):
+        return self._binop(other, "mulf")
+
+    def __truediv__(self, other):
+        return self._binop(other, "divf")
 
     def sum(self) -> "F16Vector":
         """Reduce to vector<1xf16> via tiny.sum."""
         result_type = VectorType.get([1], F16Type.get())
         op = Operation.create(
             "tiny.sum",
             results=[result_type],
-            operands=[self._value],
+            operands=[self],
         )
-        return F16Vector._wrap(op.result)
+        return op.result
+
+    def __repr__(self):
+        return repr(self).replace("Value", "F16Vector")
+
+
+@register_value_caster(VectorType.static_typeid)
+def maybe_wrap_vector(val: Value):
+    if isinstance(val.type.element_type, F16Type):
+        return F16Vector(val)
+    return val
 
 
 class Ptr:
     """Wrapper for !tiny.ptr SSA values."""
+
     _value: Value
 
     @staticmethod
@@ -113,30 +134,32 @@ def load(self, offset: Index, num_elements: int) -> F16Vector:
         op = Operation.create(
             "tiny.load",
             results=[vec_type],
-            operands=[self._value, offset._value],
+            operands=[self._value, offset],
         )
-        return F16Vector._wrap(op.result)
+        return op.result
 
     def store(self, offset: Index, vec: F16Vector) -> None:
         """Store vector<Nxf16> to pointer at offset."""
         Operation.create(
             "tiny.store",
             results=[],
-            operands=[vec._value, self._value, offset._value],
+            operands=[vec, self._value, offset],
         )
 
 
 # --- Convenience functions ---
 
 from ..compiler import MLIRModule, TutorialOpt
 
+
 def _get_type_map():
     """Type map for ch1 tiny dialect."""
     return {
         Ptr: (Ptr.get_type, Ptr),
         Index: (IndexType.get, Index),
     }
 
+
 def print_ir(fn):
     """Print generated Tiny dialect IR (verified and pretty-printed)."""
     opt = TutorialOpt()
@@ -145,6 +168,7 @@ def print_ir(fn):
         print(tiny_ir)
     return fn
 
+
 def compile_and_print(fn):
     """Compile and print all lowering stages."""
     opt = TutorialOpt()
@@ -159,7 +183,10 @@ def compile_and_print(fn):
     print("=== After tiny-to-arith ===")
     print(arith_ir)
 
-    llvm_ir = opt.run(tiny_ir, ["tiny-to-arith", "canonicalize", "cse", "tiny-to-llvm", "convert-to-llvm"])
+    llvm_ir = opt.run(
+        tiny_ir,
+        ["tiny-to-arith", "canonicalize", "cse", "tiny-to-llvm", "convert-to-llvm"],
+    )
     print("=== LLVM Dialect ===")
     print(llvm_ir)
 

python/tiny/ch2/__init__.py

@@ -2,9 +2,3 @@
 
 from ..ch1 import Index, F16Vector, Ptr
 from .dsl import accumulate, print_ir, compile_and_print
-
-__all__ = [
-    "Index", "F16Vector", "Ptr",
-    "accumulate",
-    "print_ir", "compile_and_print",
-]