feat: impl high-rank primitives and their op

examples/tarai.py

@@ -1,6 +1,6 @@
 def tarai(x: int, y: int, z: int) -> int:
     if x > y:
-        return tarai( tarai(x - 1, y, z), tarai(y - 1, z, x), tarai(z - 1, x, y) )
+        return tarai(tarai(x - 1, y, z), tarai(y - 1, z, x), tarai(z - 1, x, y))
     else:
         return y
 

examples/tarai_prim.py

@@ -7,7 +7,7 @@
 @native(gc="none")
 def tarai(x: Int[8], y: Int[8], z: Int[8]) -> Int[8]:
     if x > y:
-        return tarai( tarai(x - p1, y, z), tarai(y - p1, z, x), tarai(z - p1, x, y) )
+        return tarai(tarai(x - p1, y, z), tarai(y - p1, z, x), tarai(z - p1, x, y))
     else:
         return y
 

examples/tensor_prim.py

@@ -1,73 +1,57 @@
 from lyrt import from_prim
-from lyrt.prim import Float, Matrix, Tensor, Vector
+from lyrt.prim import Float, Matrix, Tensor
 
-# ゼロ初期化
+# 演算テスト
 
-v = Vector[Float[32], 4].zeros()
-m = Matrix[Float[32], 3, 4].zeros()
-t = Tensor[Float[32], 2, 3, 4].zeros()
-
-print(from_prim(v))
-print(from_prim(m))
-print(from_prim(t))
-
-# 初期化
-
-v2 = Vector[Float[32], 4]([0.3, -1.2, 4.5, 2.1])
-
-m2 = Matrix[Float[32], 3, 4](
+m4 = Matrix[Float[32], 2, 3](
     [
-        [1.1, -0.7, 3.3, 0.0],
-        [2.4, 5.6, -2.2, 1.9],
-        [0.5, 4.8, -1.1, 3.7],
+        [1.0, 2.0, 3.0],
+        [4.0, 5.0, 6.0],
     ]
 )
-
-t2 = Tensor[Float[32], 2, 3, 4](
+m5 = Matrix[Float[32], 2, 3](
     [
-        [
-            [0.9, -1.3, 2.2, 4.4],
-            [3.1, 0.8, -0.5, 6.6],
-            [7.7, -2.4, 1.2, 0.3],
-        ],
-        [
-            [5.5, 2.6, -3.3, 1.4],
-            [8.8, -0.9, 4.0, 2.2],
-            [6.1, 3.3, -1.7, 9.9],
-        ],
+        [0.5, 1.5, 2.5],
+        [3.5, 4.5, 5.5],
     ]
 )
 
-print(from_prim(v2))
-print(from_prim(m2))
-print(from_prim(t2))
-
-# 次元数を推論
-
-v3 = Vector[Float[32], ...]([1.0, 2.0, 3.0, 4.0, 5.0])
+print(from_prim(m4 + m5))
+print(from_prim(m4 - m5))
+print(from_prim(m4 * m5))
+print(
+    from_prim(
+        m4
+        @ Matrix[Float[32], 3, 2](
+            [
+                [1.0, 0.0],
+                [0.0, 1.0],
+                [1.0, 1.0],
+            ]
+        )
+    )
+)
 
-m3 = Matrix[Float[32], ...](
+t4 = Tensor[Float[32], 2, 2](
     [
-        [1.0, 2.0, 3.0],
-        [4.0, 5.0, 6.0],
+        [1.0, 2.0],
+        [3.0, 4.0],
     ]
 )
-
-t3 = Tensor[Float[32], ...](
+t5 = Tensor[Float[32], 2, 2](
     [
-        [
-            [1.0, 2.0],
-            [3.0, 4.0],
-        ],
+        [0.1, 0.2],
+        [0.3, 0.4],
     ]
 )
 
-print(from_prim(v3))
-print(from_prim(m3))
-print(from_prim(t3))
+print(from_prim(t4 + t5))
+print(from_prim(t4 - t5))
+print(from_prim(t4 * t5))
 
 # 0階テンソルはスカラーと同値
-t4 = Tensor[Float[32], ...](3.14)
-assert from_prim(t4) == Float[32](3.14)
-
-print(from_prim(t4))
+#
+# t4 = Tensor[Float[32], ...](3.14)
+# assert from_prim(t4) == Float[32](3.14)
+#
+# print(from_prim(t4))

pyproject.toml

@@ -4,7 +4,10 @@ version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
-dependencies = []
+dependencies = [
+    "numpy>=2.4.1",
+    "pyyaml>=6.0.3",
+]
 
 [dependency-groups]
 dev = [

src/lyrt/__init__.pyi

@@ -2,9 +2,19 @@ from typing import Callable, Literal, TypeVar
 
 from . import prim
 
-__all__ = ["prim", "native", "to_prim", "from_prim"]
+__all__ = ["prim", "native", "to_prim", "from_prim", "alloc", "dealloc"]
 
 T = TypeVar("T")
+AllocT = TypeVar(
+    "AllocT",
+    bound=(
+        prim.Prim[prim.Int]
+        | prim.Prim[prim.Float]
+        | prim.Vector
+        | prim.Matrix
+        | prim.Tensor
+    ),
+)
 PrimT = TypeVar("PrimT", bound=prim.Prim[prim.Int] | prim.Prim[prim.Float])
 PrimFunc = TypeVar(
     "PrimFunc",
@@ -25,3 +35,13 @@ def from_prim(
         | prim.Tensor
     ),
 ) -> object: ...
+def alloc(value: AllocT) -> AllocT: ...
+def dealloc(
+    value: (
+        prim.Prim[prim.Int]
+        | prim.Prim[prim.Float]
+        | prim.Vector
+        | prim.Matrix
+        | prim.Tensor
+    ),
+) -> None: ...

src/lyrt/prim/__init__.pyi

@@ -61,6 +61,9 @@ Shape = TypeVarTuple("Shape")
 class Vector:
     def __class_getitem__(cls, key: tuple[Prim[Int | Float], int]) -> Type[Vector]: ...
     def __init__(self, value: List[NumberLike]) -> None: ...
+    def __add__(self, other: Vector) -> Vector: ...
+    def __sub__(self, other: Vector) -> Vector: ...
+    def __mul__(self, other: Vector) -> Vector: ...
     @classmethod
     def zeros(cls) -> Vector: ...
     @classmethod
@@ -73,6 +76,10 @@ class Matrix:
         cls, key: tuple[Prim[Int | Float], int, int]
     ) -> Type[Matrix]: ...
     def __init__(self, value: List[List[NumberLike]]) -> None: ...
+    def __add__(self, other: Matrix) -> Matrix: ...
+    def __sub__(self, other: Matrix) -> Matrix: ...
+    def __mul__(self, other: Matrix) -> Matrix: ...
+    def __matmul__(self, other: Matrix) -> Matrix: ...
     @classmethod
     def zeros(cls) -> Matrix: ...
     @classmethod
@@ -87,6 +94,9 @@ class Tensor:
 
     type _NestedNumberLike = NumberLike | List[_NestedNumberLike]
     def __init__(self, value: _NestedNumberLike) -> None: ...
+    def __add__(self, other: Tensor) -> Tensor: ...
+    def __sub__(self, other: Tensor) -> Tensor: ...
+    def __mul__(self, other: Tensor) -> Tensor: ...
     @classmethod
     def zeros(cls) -> Tensor: ...
     @classmethod

src/lython/dialects/cpp/PyVerifier.cpp

@@ -446,6 +446,12 @@ LogicalResult CastFromPrimOp::verify() {
       return success();
   }
 
+  // Allow ranked tensor types to !py.str conversion (repr carrier)
+  if (llvm::isa<::mlir::RankedTensorType>(inputType)) {
+    if (resultType == StrType::get(ctx))
+      return success();
+  }
+
   return emitOpError("unsupported type conversion from ")
          << inputType << " to " << resultType;
 }

src/lython/lowering/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_library(LythonRuntimeLowering STATIC
     RuntimeSupport.cpp
     RuntimeLoweringPass.cpp
+    LinalgLoweringPass.cpp
     NativeVerificationPass.cpp
     PyFuncLowering.cpp
     PyOptimizationPass.cpp
@@ -23,7 +24,14 @@ target_link_libraries(LythonRuntimeLowering
     PRIVATE
     PyDialect
     MLIRIR
+    MLIRArithDialect
+    MLIRControlFlowDialect
     MLIRFuncDialect
     MLIRLLVMDialect
+    MLIRLinalgDialect
+    MLIRLinalgToStandard
+    MLIRMemRefDialect
+    MLIRSCFDialect
     MLIRSupport
+    MLIRTensorDialect
 )

src/lython/lowering/LinalgLoweringPass.cpp

@@ -0,0 +1,53 @@
+// This pass lowers linalg ops to standard/scf while allowing other dialects.
+
+#include "RuntimeSupport.h"
+
+#include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace py {
+namespace {
+
+struct LinalgLoweringPass
+    : public PassWrapper<LinalgLoweringPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LinalgLoweringPass)
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    MLIRContext *ctx = module.getContext();
+    RewritePatternSet patterns(ctx);
+    linalg::populateLinalgToStandardConversionPatterns(patterns);
+
+    ConversionTarget target(*ctx);
+    target.addLegalOp<ModuleOp>();
+    target.addLegalDialect<arith::ArithDialect, func::FuncDialect,
+                           tensor::TensorDialect, scf::SCFDialect,
+                           memref::MemRefDialect, cf::ControlFlowDialect,
+                           LLVM::LLVMDialect>();
+    target.addIllegalDialect<linalg::LinalgDialect>();
+    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+
+    if (failed(applyPartialConversion(module, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createLinalgLoweringPass() {
+  return std::make_unique<LinalgLoweringPass>();
+}
+
+} // namespace py