[nfc] Automate OpTileLowering and pattern registration using .def file (part of #132) (#180)

zoecarver · web-flow · commit ba412e1f32cb · 2026-01-06T10:21:40.000-05:00
* generate patterns from .def * add docs * not yet: automate defs in TTLops.td Does not use yaml file as suggested in #132 as that seemed overkill for this change, may still be good to do in the future.
diff --git a/docs/sphinx/contributor-guide.md b/docs/sphinx/contributor-guide.md
@@ -14,3 +14,47 @@
 - Add new user-facing pages under `docs/sphinx` and link them in `index.rst`.
 - Keep contributor-only instructions in this guide or `guidelines.md`.
 - Build docs with `cmake --build build --target ttlang-docs`.
+
+## Adding Elementwise Operations
+
+To add a new elementwise operation (unary or binary), update these files:
+
+### 1. `include/ttlang/Dialect/TTL/TTLElementwiseOps.def`
+
+Add an entry with the TTL op name, tile op name, and TTKernel init/compute op names:
+
+```cpp
+// Binary op (3-arg form: DST[odst] = op(DST[src0], DST[src1]))
+TTL_BINARY_TILE_OP(NewOp, NewOpTileOp, NewOpBinaryTilesInitOp, NewOpBinaryTilesOp)
+
+// Unary op (in-place form: DST[dst_idx] = op(DST[dst_idx]))
+TTL_UNARY_TILE_OP(NewOp, NewOpTileOp, NewOpTileInitOp, NewOpTileOp)
+
+// Special binary op (2-arg in-place form, like Max)
+TTL_BINARY_TILE_OP_SPECIAL(NewOp, NewOpTileOp, NewOpTilesInitOp, NewOpTilesOp)
+```
+
+This automatically generates:
+- C++ lowering patterns (`ConvertTTLToCompute.cpp`, `ConvertTTLTileOpsToTTKernel.cpp`)
+- Python bindings (`_generated_elementwise.py`)
+
+### 2. `include/ttlang/Dialect/TTL/IR/TTLOps.td`
+
+Add the TableGen op definitions using the multiclass:
+
+```tablegen
+// Binary op
+defm TTL_NewOp : TTL_BinaryElementwisePair<"newop", "newop_tiles">;
+
+// Unary op
+defm TTL_NewOp : TTL_UnaryElementwisePair<"newop", "newop_tile">;
+```
+
+### 3. Verify the TTKernel ops exist in tt-mlir
+
+The TTKernel init and compute ops must exist in `tt-mlir/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td`. If they don't, they need to be added to tt-mlir first.
+
+### 4. Add tests
+
+- Add a lit test in `test/ttlang/Dialect/TTL/Transforms/` for the lowering
+- Add to `test/python/test_elementwise_ops.py` for end-to-end verification
diff --git a/include/ttlang/Dialect/TTL/TTLElementwiseOps.def b/include/ttlang/Dialect/TTL/TTLElementwiseOps.def
@@ -6,47 +6,72 @@
 // TTL Elementwise Operation Definitions
 //===----------------------------------------------------------------------===//
 //
-// This file defines mappings for TTL elementwise operations lowering to
-// ttl.compute with tile ops.
+// This file defines mappings for TTL elementwise operations, supporting:
+// 1. TTL tensor ops -> ttl.compute with tile ops (ConvertTTLToCompute)
+// 2. TTL tile ops -> TTKernel ops (ConvertTTLTileOpsToTTKernel)
+// 3. Python binding generation (gen_elementwise.py)
 //
 // It uses the X-macro pattern to generate boilerplate code for pattern
 // definitions and registration.
 //
+// Macro signatures:
+//   TTL_BINARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)
+//   TTL_UNARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)
+//   TTL_BINARY_TILE_OP_SPECIAL(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)
+//
+// Parameters:
+//   TTL_OP     - TTL tensor op name (e.g., Add -> ttl::AddOp)
+//   TILE_OP   - TTL tile op name (e.g., AddTileOp -> ttl::AddTileOp)
+//   TTK_INIT  - TTKernel init op (e.g., AddBinaryTilesInitOp)
+//   TTK_COMPUTE - TTKernel compute op (e.g., AddBinaryTilesOp)
+//
 // Usage:
-//   #define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP)  // your code here
-//   #define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP)   // your code here
+//   #define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE) ...
+//   #define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE) ...
 //   #include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
 //
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
-// Elementwise operation mappings (TTL tensor ops -> ttl.compute with tile ops)
+// Elementwise operation mappings
 //===----------------------------------------------------------------------===//
 
 #ifndef TTL_BINARY_TILE_OP
-#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP)
+#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)
 #endif
 
 #ifndef TTL_UNARY_TILE_OP
-#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP)
+#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)
+#endif
+
+// Special macro for binary ops that use a different lowering template.
+// Max uses 2-arg in-place form: DST[dst0] = max(DST[dst0], DST[dst1])
+#ifndef TTL_BINARY_TILE_OP_SPECIAL
+#define TTL_BINARY_TILE_OP_SPECIAL(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)
 #endif
 
-// Binary operations: TTL tensor op -> TTL tile op
-TTL_BINARY_TILE_OP(Add, AddTileOp)
-TTL_BINARY_TILE_OP(Sub, SubTileOp)
-TTL_BINARY_TILE_OP(Mul, MulTileOp)
-TTL_BINARY_TILE_OP(Max, MaxTileOp)
-
-// Unary operations: TTL tensor op -> TTL tile op
-TTL_UNARY_TILE_OP(Exp, ExpTileOp)
-TTL_UNARY_TILE_OP(Log, LogTileOp)
-TTL_UNARY_TILE_OP(Sqrt, SqrtTileOp)
-TTL_UNARY_TILE_OP(Rsqrt, RsqrtTileOp)
-TTL_UNARY_TILE_OP(Tanh, TanhTileOp)
-TTL_UNARY_TILE_OP(Abs, AbsTileOp)
-TTL_UNARY_TILE_OP(Neg, NegTileOp)
-TTL_UNARY_TILE_OP(Relu, ReluTileOp)
-TTL_UNARY_TILE_OP(Sigmoid, SigmoidTileOp)
+// Binary operations: TTL tensor op -> TTL tile op -> TTKernel init/compute ops
+// These use the standard 3-arg binary template: DST[odst] = op(DST[src0], DST[src1])
+TTL_BINARY_TILE_OP(Add, AddTileOp, AddBinaryTilesInitOp, AddBinaryTilesOp)
+TTL_BINARY_TILE_OP(Sub, SubTileOp, SubBinaryTilesInitOp, SubBinaryTilesOp)
+TTL_BINARY_TILE_OP(Mul, MulTileOp, MulBinaryTilesInitOp, MulBinaryTilesOp)
+
+// Special binary ops with non-standard lowering
+// Max uses 2-arg in-place form (TTLTileMaxToTTKernel template)
+TTL_BINARY_TILE_OP_SPECIAL(Max, MaxTileOp, BinaryMaxTileInitOp, BinaryMaxTileOp)
+
+// Unary operations: TTL tensor op -> TTL tile op -> TTKernel init/compute ops
+// These use the standard unary template: DST[dst_idx] = op(DST[dst_idx])
+TTL_UNARY_TILE_OP(Exp, ExpTileOp, ExpTileInitOp, ExpTileOp)
+TTL_UNARY_TILE_OP(Log, LogTileOp, LogTileInitOp, LogTileOp)
+TTL_UNARY_TILE_OP(Sqrt, SqrtTileOp, SqrtTileInitOp, SqrtTileOp)
+TTL_UNARY_TILE_OP(Rsqrt, RsqrtTileOp, RsqrtTileInitOp, RsqrtTileOp)
+TTL_UNARY_TILE_OP(Tanh, TanhTileOp, TanhTileInitOp, TanhTileOp)
+TTL_UNARY_TILE_OP(Abs, AbsTileOp, AbsTileInitOp, AbsTileOp)
+TTL_UNARY_TILE_OP(Neg, NegTileOp, NegativeTileInitOp, NegativeTileOp)
+TTL_UNARY_TILE_OP(Relu, ReluTileOp, ReluTileInitOp, ReluTileOp)
+TTL_UNARY_TILE_OP(Sigmoid, SigmoidTileOp, SigmoidTileInitOp, SigmoidTileOp)
 
 #undef TTL_BINARY_TILE_OP
 #undef TTL_UNARY_TILE_OP
+#undef TTL_BINARY_TILE_OP_SPECIAL
diff --git a/lib/Dialect/TTL/Transforms/ConvertTTLTileOpsToTTKernel.cpp b/lib/Dialect/TTL/Transforms/ConvertTTLTileOpsToTTKernel.cpp
@@ -353,46 +353,26 @@ struct TTLTileCopyToTTKernel : OpConversionPattern<CopyTileOp> {
 };
 
 //===----------------------------------------------------------------------===//
-// Unary Tile Op Lowerings (LLVM-style type aliases)
+// Tile Op Lowerings - Generated from TTLElementwiseOps.def
 //===----------------------------------------------------------------------===//
 
-using ExpTileLowering =
-    TTLTileUnaryToTTKernel<ExpTileOp, ttk::ExpTileInitOp, ttk::ExpTileOp>;
-using LogTileLowering =
-    TTLTileUnaryToTTKernel<LogTileOp, ttk::LogTileInitOp, ttk::LogTileOp>;
-using SqrtTileLowering =
-    TTLTileUnaryToTTKernel<SqrtTileOp, ttk::SqrtTileInitOp, ttk::SqrtTileOp>;
-using RsqrtTileLowering =
-    TTLTileUnaryToTTKernel<RsqrtTileOp, ttk::RsqrtTileInitOp, ttk::RsqrtTileOp>;
-using TanhTileLowering =
-    TTLTileUnaryToTTKernel<TanhTileOp, ttk::TanhTileInitOp, ttk::TanhTileOp>;
-using SigmoidTileLowering =
-    TTLTileUnaryToTTKernel<SigmoidTileOp, ttk::SigmoidTileInitOp,
-                           ttk::SigmoidTileOp>;
-using AbsTileLowering =
-    TTLTileUnaryToTTKernel<AbsTileOp, ttk::AbsTileInitOp, ttk::AbsTileOp>;
-using NegTileLowering =
-    TTLTileUnaryToTTKernel<NegTileOp, ttk::NegativeTileInitOp,
-                           ttk::NegativeTileOp>;
-using ReluTileLowering =
-    TTLTileUnaryToTTKernel<ReluTileOp, ttk::ReluTileInitOp, ttk::ReluTileOp>;
+// Generate type aliases for unary tile op lowerings
+#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)              \
+  using TTL_OP##TileLowering =                                                 \
+      TTLTileUnaryToTTKernel<TILE_OP, ttk::TTK_INIT, ttk::TTK_COMPUTE>;
+#include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
 
-//===----------------------------------------------------------------------===//
-// Binary Tile Op Lowerings
-//===----------------------------------------------------------------------===//
+// Generate type aliases for binary tile op lowerings (standard 3-arg form)
+#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)             \
+  using TTL_OP##TileLowering =                                                 \
+      TTLTileBinaryToTTKernel<TILE_OP, ttk::TTK_INIT, ttk::TTK_COMPUTE>;
+#include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
 
-using AddTileLowering =
-    TTLTileBinaryToTTKernel<AddTileOp, ttk::AddBinaryTilesInitOp,
-                            ttk::AddBinaryTilesOp>;
-using SubTileLowering =
-    TTLTileBinaryToTTKernel<SubTileOp, ttk::SubBinaryTilesInitOp,
-                            ttk::SubBinaryTilesOp>;
-using MulTileLowering =
-    TTLTileBinaryToTTKernel<MulTileOp, ttk::MulBinaryTilesInitOp,
-                            ttk::MulBinaryTilesOp>;
-using MaxTileLowering =
-    TTLTileMaxToTTKernel<MaxTileOp, ttk::BinaryMaxTileInitOp,
-                         ttk::BinaryMaxTileOp>;
+// Generate type aliases for special binary tile op lowerings (2-arg in-place)
+#define TTL_BINARY_TILE_OP_SPECIAL(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)     \
+  using TTL_OP##TileLowering =                                                 \
+      TTLTileMaxToTTKernel<TILE_OP, ttk::TTK_INIT, ttk::TTK_COMPUTE>;
+#include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
 
 } // namespace
 
@@ -409,14 +389,21 @@ void populateTTLTileOpsToTTKernelPatterns(TypeConverter *typeConverter,
   patterns.add<TTLTileRegsAcquireToTTKernel, TTLTileRegsCommitToTTKernel,
                TTLTileRegsWaitToTTKernel, TTLTileRegsReleaseToTTKernel>(ctx);
 
-  // Tile op lowerings (ttl.tile_* → ttkernel.*_tile)
-  patterns.add<
-      // Unary ops
-      ExpTileLowering, LogTileLowering, SqrtTileLowering, RsqrtTileLowering,
-      TanhTileLowering, SigmoidTileLowering, AbsTileLowering, NegTileLowering,
-      ReluTileLowering,
-      // Binary ops
-      AddTileLowering, SubTileLowering, MulTileLowering, MaxTileLowering>(ctx);
+  // Tile op lowerings - generated from TTLElementwiseOps.def
+  // Unary ops (ttl.tile_* → ttkernel.*_tile)
+#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)              \
+  patterns.add<TTL_OP##TileLowering>(ctx);
+#include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
+
+  // Binary ops (ttl.tile_* → ttkernel.*_tiles)
+#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)             \
+  patterns.add<TTL_OP##TileLowering>(ctx);
+#include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
+
+  // Special binary ops (non-standard lowering template)
+#define TTL_BINARY_TILE_OP_SPECIAL(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)     \
+  patterns.add<TTL_OP##TileLowering>(ctx);
+#include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
 
   // Copy op needs the type converter.
   patterns.add<TTLTileCopyToTTKernel>(*typeConverter, ctx);
diff --git a/lib/Dialect/TTL/Transforms/ConvertTTLToCompute.cpp b/lib/Dialect/TTL/Transforms/ConvertTTLToCompute.cpp
@@ -281,10 +281,13 @@ struct LowerUnaryToCompute : OpRewritePattern<TTLOp> {
 //===----------------------------------------------------------------------===//
 
 // Generate type aliases for binary operations using tile ops
-#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP)                                    \
+// (TTK_INIT and TTK_COMPUTE are unused here, only needed for TTKernel lowering)
+#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)             \
+  using Lower##TTL_OP = LowerBinaryToCompute<TTL_OP##Op, TILE_OP>;
+#define TTL_BINARY_TILE_OP_SPECIAL(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)     \
   using Lower##TTL_OP = LowerBinaryToCompute<TTL_OP##Op, TILE_OP>;
 // Generate type aliases for unary operations using tile ops
-#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP)                                     \
+#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)              \
   using Lower##TTL_OP = LowerUnaryToCompute<TTL_OP##Op, TILE_OP>;
 #include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
 
@@ -319,8 +322,14 @@ void populateTTLToComputePatterns(RewritePatternSet &patterns) {
 
   // Register patterns for lowering to ttl.compute with tile ops.
   // These are generated from TTLElementwiseOps.def using tile-based mappings.
-#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP) patterns.add<Lower##TTL_OP>(ctx);
-#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP) patterns.add<Lower##TTL_OP>(ctx);
+  // (TTK_INIT and TTK_COMPUTE are unused here, only needed for TTKernel
+  // lowering)
+#define TTL_BINARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)             \
+  patterns.add<Lower##TTL_OP>(ctx);
+#define TTL_BINARY_TILE_OP_SPECIAL(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)     \
+  patterns.add<Lower##TTL_OP>(ctx);
+#define TTL_UNARY_TILE_OP(TTL_OP, TILE_OP, TTK_INIT, TTK_COMPUTE)              \
+  patterns.add<Lower##TTL_OP>(ctx);
 #include "ttlang/Dialect/TTL/TTLElementwiseOps.def"
 }
 
diff --git a/python/gen_elementwise.py b/python/gen_elementwise.py
@@ -64,24 +64,43 @@ def {name}(input: "TensorBlock") -> "TensorBlock":
 
 
 def parse_def_file(def_path: Path) -> tuple[list[str], list[str]]:
-    """Parse TTLElementwiseOps.def and extract operation names."""
+    """Parse TTLElementwiseOps.def and extract operation names.
+
+    Handles both standard and special binary ops:
+    - TTL_BINARY_TILE_OP(Name, TileOp, TTKInit, TTKCompute)
+    - TTL_BINARY_TILE_OP_SPECIAL(Name, TileOp, TTKInit, TTKCompute)
+    - TTL_UNARY_TILE_OP(Name, TileOp, TTKInit, TTKCompute)
+    """
     content = def_path.read_text()
 
     binary_ops = []
     unary_ops = []
 
-    # Match TTL_BINARY_TILE_OP(Name, TileOp) but skip #define lines
-    for match in re.finditer(r"^TTL_BINARY_TILE_OP\((\w+),", content, re.MULTILINE):
+    # Match TTL_BINARY_TILE_OP(Name, ...) and TTL_BINARY_TILE_OP_SPECIAL(Name, ...)
+    # but skip #define lines
+    for match in re.finditer(
+        r"^TTL_BINARY_TILE_OP(?:_SPECIAL)?\((\w+),", content, re.MULTILINE
+    ):
         name = match.group(1).lower()
         # Skip macro parameter names (lowercase indicates it's a parameter)
-        if name[0].isupper() or name not in ("ttl_op", "tile_op"):
+        if name[0].isupper() or name not in (
+            "ttl_op",
+            "tile_op",
+            "ttk_init",
+            "ttk_compute",
+        ):
             binary_ops.append(name)
 
-    # Match TTL_UNARY_TILE_OP(Name, TileOp) but skip #define lines
+    # Match TTL_UNARY_TILE_OP(Name, ...) but skip #define lines
     for match in re.finditer(r"^TTL_UNARY_TILE_OP\((\w+),", content, re.MULTILINE):
         name = match.group(1).lower()
         # Skip macro parameter names
-        if name[0].isupper() or name not in ("ttl_op", "tile_op"):
+        if name[0].isupper() or name not in (
+            "ttl_op",
+            "tile_op",
+            "ttk_init",
+            "ttk_compute",
+        ):
             unary_ops.append(name)
 
     return binary_ops, unary_ops
