Feat: Tensor, DType, PatternMatcher, etc by hikettei · Pull Request #21 · hikettei/tiny_polyhedral_compiler

hikettei · 2025-12-24T18:25:32Z

github-actions · 2025-12-24T18:25:42Z

🤖 Hi @hikettei, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

gemini-code-assist · 2025-12-24T18:25:49Z

Summary of Changes

Hello @hikettei, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request lays the foundational architecture for a new tensor operations framework, focusing on establishing a robust Intermediate Representation (IR) and a flexible backend abstraction. It introduces structured data type definitions, a comprehensive set of IR operations, and an initial CPU backend. The changes are a significant step towards enabling advanced features like polyhedral graph generation and deterministic loop fusion, as indicated by the 'WIP' status and the checklist in the description.

Highlights

Intermediate Representation (IR) Foundation: Introduced caten/ir.py which defines core components for an IR, including ATenAxis, ATenOpType, and ATenOp. This file also establishes various unary, binary, and ternary operations, allocation mechanisms, JIT-related operations, and control flow constructs like Range, Loop, and When.
Data Type Definitions: Added caten/dtype.py to manage data types. It includes a DType class with a metaclass for caching, and pre-defines common types such as float64, float32, int64, and int32.
Tensor API and Backend Abstraction: Created caten/tensor.py which defines the high-level ATenSpec and ATen classes. It also introduces ATenBase as an abstract base class for tensor operations, providing a mechanism to register and dispatch to different hardware backends. The Tensor class now acts as a factory, instantiating the appropriate backend implementation.
CPU Backend Implementation: Implemented a basic CPU backend in caten/runtime/cpu.py with the CPUTensor class, which registers itself with the ATenBase system. This provides a concrete target for the tensor operations.
Cleanup and Testing: Removed the caten/ops.py file, indicating a refactoring and replacement of its previous, less structured content. A new test file, test/test_kernel.py, was added with a basic test case to instantiate a C.Tensor.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

github-actions · 2025-12-24T18:26:35Z

🤖 I'm sorry @hikettei, but I was unable to process your request. Please see the logs for more details.

gemini-code-assist

Code Review

This pull request lays the groundwork for a new tensor library, including data types, an intermediate representation (IR), and a backend dispatch system. While it's clearly a work-in-progress, there are several critical issues such as missing imports, undefined name errors, and incorrect class definitions that will prevent the code from running. My review focuses on addressing these fundamental correctness problems to make the code executable, and also includes suggestions to improve code structure and adhere to Python conventions.

gemini-code-assist · 2025-12-24T18:28:08Z

caten/tensor.py

+class ATen:
+    op: ATenOp # ATen is just a wrapper for ATenOp


ATen is instantiated as if it were a dataclass in the apply method, but it's not defined as one. This will cause a TypeError. You should decorate it with @dataclass. Additionally, ATenOp is not defined in this scope and should be referenced as ir.ATenOp.

Suggested change

class ATen:

op: ATenOp # ATen is just a wrapper for ATenOp

@dataclass(frozen=True)

class ATen:

op: ir.ATenOp # ATen is just a wrapper for ATenOp

gemini-code-assist · 2025-12-24T18:28:08Z

caten/tensor.py

+class ATen:
+    op: ATenOp # ATen is just a wrapper for ATenOp
+    @classmethod
+    def from_shape(cls, shape: List[ATenOp]):


ATenOp is not defined in this scope. It should be referenced from the ir module as ir.ATenOp.

Suggested change

def from_shape(cls, shape: List[ATenOp]):

def from_shape(cls, shape: List[ir.ATenOp]):

gemini-code-assist · 2025-12-24T18:28:08Z

caten/tensor.py

+    def from_shape(cls, shape: List[ATenOp]):
+        return ir.Allocate(shape) # TODO
+
+    def apply(self, op: Callable, *args: List, **kwargs) -> ATen: return ATen(op=op(*args, **kwargs))


The type hint *args: List is invalid syntax. To type variadic positional arguments, you can use *args: Any. It's also good practice to type **kwargs as **kwargs: Any.

Suggested change

def apply(self, op: Callable, *args: List, **kwargs) -> ATen: return ATen(op=op(*args, **kwargs))

def apply(self, op: Callable, *args: Any, **kwargs: Any) -> ATen: return ATen(op=op(*args, **kwargs))

gemini-code-assist · 2025-12-24T18:28:08Z

caten/tensor.py

+from typing import Any, Optional, Tuple, Union, ClassVar
+import os


Some necessary imports are missing. List and Callable should be imported from typing, and dataclass should be imported from dataclasses to be used as a decorator.

Suggested change

from typing import Any, Optional, Tuple, Union, ClassVar

import os

from typing import Any, Optional, Tuple, Union, ClassVar, List, Callable

from dataclasses import dataclass

import os

gemini-code-assist · 2025-12-24T18:28:08Z

caten/ir.py

+
+@dataclass(frozen=True)
+class ATenOp(metaclass=ABCMeta):
+    args: List[AtenOp]


There's a typo in the type hint. AtenOp should be ATenOp to match the class definition.

Suggested change

args: List[AtenOp]

args: List[ATenOp]

gemini-code-assist · 2025-12-24T18:28:08Z

caten/dtype.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+from dataclasses import dataclass, fields


The fields function is imported from dataclasses but is not used in this file. It's good practice to remove unused imports to keep the code clean.

Suggested change

from dataclasses import dataclass, fields

from dataclasses import dataclass

gemini-code-assist · 2025-12-24T18:28:08Z

caten/__init__.py

+__all__ = [
+    "dtype",
+    "helpers",
+    "ir",
+    "tensor"
+]


The __all__ list is missing key public APIs like Tensor. To provide a better user experience, it's recommended to export the main classes and functions that users are expected to use directly. This makes them discoverable and allows for imports like from caten import Tensor.

Suggested change

__all__ = [

"dtype",

"helpers",

"ir",

"tensor"

]

__all__ = [

"dtype",

"helpers",

"ir",

"tensor",

"Tensor",

"get_backend",

"ATenSpec"

]

gemini-code-assist · 2025-12-24T18:28:08Z

caten/ir.py

+def Var():
+    pass


The function Var is defined but has an empty body. This should be implemented or removed if it's a placeholder.

gemini-code-assist · 2025-12-24T18:28:08Z

caten/tensor.py

+        impl = DEVICE_TO_TENSOR.get(get_backend())
+        if impl is None: raise ValueError(f"Unknown BACKEND={get_backend()}")


For efficiency and clarity, you can store the result of get_backend() in a variable to avoid calling it twice.

Suggested change

impl = DEVICE_TO_TENSOR.get(get_backend())

if impl is None: raise ValueError(f"Unknown BACKEND={get_backend()}")

backend = get_backend()

impl = DEVICE_TO_TENSOR.get(backend)

if impl is None: raise ValueError(f"Unknown BACKEND={backend}")

gemini-code-assist · 2025-12-24T18:28:08Z

caten/tensor.py

+class range():
+    pass
+
+class when():
+    pass


According to PEP 8, class names should use PascalCase. Please rename range to Range and when to When. Renaming range also avoids shadowing the built-in range function, which can be a source of confusion.

Suggested change

class range():

pass

class when():

pass

class Range():

pass

class When():

pass

hikettei added 5 commits December 25, 2025 02:21

wip

101b88b

concepts

9578c37

wip

fec84c0

wip

7856871

wip

2d4c0ce

hikettei changed the title ~~WIP: TensorOps~~ WIP: PyTorch/Triton Style Frontend Dec 24, 2025

gemini-code-assist bot reviewed Dec 24, 2025

View reviewed changes

hikettei added 20 commits December 25, 2025 03:28

facet ideas

aa350ec

allocation graph

d800a8e

allocation graph

51de548

allocation graph

557dc10

allocation graph

2ab506b

concrete semantics

673906c

dataclass

e4ddd5e

dataclass

3ac5d96

dataclass

e66e626

dataclass

e4bacfd

tuple

a3056eb

concrete semantics

e259924

concrete semantics

227b779

concrete semantics

19db18b

WIP

74471c9

some clean ups

a5c0b82

just an idea

7e60cc0

feat: reshape

71fc714

feat: reshape

704b6be

feat: expand

0b526ce

hikettei added 17 commits December 25, 2025 19:19

Cache, Type Inference, Verification

bc0f4a1

Cache, Type Inference, Verification

7470551

Cache, Type Inference, Verification

d66007c

Feat: Simplifier

8d68192

Feat: Simplifier

5820b67

Feat: Simplifier

fa3ae24

Reshape symbolic

d43c14b

docs

6437e69

docs

f35442a

smax

fc442a6

write specs on def

3061ab2

write specs on def

d6e83da

Fix scope issue

10b0a41

Fix scope issue

b0a7749

fix ruff

c34c90a

fix ruff

51abfad

fix ruff

6215735

hikettei changed the title ~~WIP: PyTorch/Triton Style Frontend~~ Feat: Tensor, DType, PatternMatcher, etc Dec 25, 2025

hikettei marked this pull request as ready for review December 25, 2025 13:10

hikettei added 9 commits December 25, 2025 22:18

helpers

e72883f

helpers

50372b0

helpers

77c9157

helpers

4700274

helpers

a2ce359

helpers

dec32ae

helpers

8c5e351

helpers

8730f7c

helpers

1cb1d46

hikettei merged commit f320079 into main Dec 25, 2025
4 checks passed

hikettei deleted the tensor branch December 25, 2025 13:57

	def from_shape(cls, shape: List[ATenOp]):
	def from_shape(cls, shape: List[ir.ATenOp]):

	def apply(self, op: Callable, args: List, kwargs) -> ATen: return ATen(op=op(args, **kwargs))
	def apply(self, op: Callable, args: Any, kwargs: Any) -> ATen: return ATen(op=op(args, **kwargs))

		from typing import Any, Optional, Tuple, Union, ClassVar
		import os

		@@ -0,0 +1,25 @@
		from __future__ import annotations
		from dataclasses import dataclass, fields

	from dataclasses import dataclass, fields
	from dataclasses import dataclass

		impl = DEVICE_TO_TENSOR.get(get_backend())
		if impl is None: raise ValueError(f"Unknown BACKEND={get_backend()}")

-        impl = DEVICE_TO_TENSOR.get(get_backend())
-        if impl is None: raise ValueError(f"Unknown BACKEND={get_backend()}")
+        backend = get_backend()
+        impl = DEVICE_TO_TENSOR.get(backend)
+        if impl is None: raise ValueError(f"Unknown BACKEND={backend}")

Conversation

hikettei commented Dec 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Dec 24, 2025

Uh oh!

gemini-code-assist bot commented Dec 24, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions bot commented Dec 24, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

hikettei commented Dec 24, 2025 •

edited

Loading