Refactoring of TTIRTraits.cpp

- improved comments in few places
tenstorrent · Dec 26, 2024 · 34a1592 · 34a1592
1 parent f75e998
commit 34a1592
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 18 deletions.
diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIRBase.td b/include/ttmlir/Dialect/TTIR/IR/TTIRBase.td
@@ -57,8 +57,14 @@ class TTIR_Trait<string name, list<Trait> traits = []> : NativeOpTrait<name, tra
   let cppNamespace = "::mlir::tt::ttir::OpTrait";
 }
 
+// Involution is property of an operation where applying the operation twice results in the original value.
+// Example: not(not(x)) == x
 def TTIR_Involution : TTIR_Trait<"TTIRInvolution", [DestinationStyleOpInterface, TwoOperands, NoMemoryEffect]>;
+// Idempotence is property of an operation where applying the operation twice results in the same value as applying it once.
+// Example: abs(abs(x)) == abs(x)
 def TTIR_Idempotence : TTIR_Trait<"TTIRIdempotence", [DestinationStyleOpInterface, TwoOperands, NoMemoryEffect]>;
+// BinaryIdempotence is property of a binary operation where applying the operation on the same value results in the same value.
+// Example: and(x, x) == x
 def TTIR_BinaryIdempotence : TTIR_Trait<"TTIRBinaryIdempotence", [DestinationStyleOpInterface, ThreeOperands, NoMemoryEffect]>;
 
 #endif
diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIRTraits.h b/include/ttmlir/Dialect/TTIR/IR/TTIRTraits.h
@@ -6,7 +6,7 @@
 #define TTMLIR_TTMLIR_DIALECT_TTIR_TTIRTRAITS_H
 
 #include "mlir/IR/OpDefinition.h"
-#include <mlir/Support/LLVM.h>
+#include "mlir/Support/LLVM.h"
 
 namespace mlir {
 namespace tt {

diff --git a/lib/Dialect/TTIR/IR/TTIRTraits.cpp b/lib/Dialect/TTIR/IR/TTIRTraits.cpp
@@ -4,17 +4,19 @@
 
 #include "ttmlir/Dialect/TTIR/IR/TTIRTraits.h"
 
-// Op is foldable iff:
-// 1. argment and result types are the same.
-// 2. argument is defined by the same op.
+// Check if all operands and result have the same type. Function assumes op has
+// at least one operand and exactly one result.
+static bool operandAndResultSameType(mlir::Operation *op) {
+  return llvm::all_equal(op->getOperandTypes()) &&
+         op->getOperand(0).getType() == op->getResult(0).getType();
+}
+
+// If Op has TTIRInvolution trait, then it's foldable if:
+// 1. Argument and result types are the same.
+// 2. Argument is defined by the same op.
 // 3. 1) is true for the producing op of the argument.
 // op(op(T a, T r0), T r1)
 bool mlir::tt::ttir::OpTrait::impl::verifyInvolution(mlir::Operation *op) {
-  // Dependant trait of TTIRInvolution is DestionationStyleOpInterface, hence
-  // operands include the result.
-  auto operandAndResultSameType = [](Operation *op) {
-    return llvm::all_equal(op->getOperandTypes());
-  };
   if (!operandAndResultSameType(op)) {
     return false;
   }
@@ -25,17 +27,12 @@ bool mlir::tt::ttir::OpTrait::impl::verifyInvolution(mlir::Operation *op) {
   return operandAndResultSameType(producerOp);
 }
 
-// Op is foldable iff:
-// 1. argment and result types are the same.
-// 2. argument is defined by the same op.
+// If Op has TTIRIdempotence trait, then it's foldable if:
+// 1. Argument and result types are the same.
+// 2. Argument is defined by the same op.
 // 3. 1) is true for the producing op of the argument.
 // op(op(T a, T r0), T r1)
 bool mlir::tt::ttir::OpTrait::impl::verifyIdempotence(mlir::Operation *op) {
-  // Dependant trait of TTIRIdempotence is DestionationStyleOpInterface, hence
-  // operands include the result.
-  auto operandAndResultSameType = [](mlir::Operation *op) {
-    return llvm::all_equal(op->getOperandTypes());
-  };
   if (!operandAndResultSameType(op)) {
     return false;
   }
@@ -46,7 +43,7 @@ bool mlir::tt::ttir::OpTrait::impl::verifyIdempotence(mlir::Operation *op) {
   return operandAndResultSameType(producerOp);
 }
 
-// Op is foldable iff:
+// If Op has TTIRBinaryIdempotence trait, then it's foldable if:
 // 1. Both inputs are the same.
 // 2. Inputs and result types are the same.
 bool mlir::tt::ttir::OpTrait::impl::verifyBinaryIdempotence(