[BACKEND] Rewrite tablegen combine into cpp patterns (#7709)

<!---
The core Triton is a small number of people, and we receive many PRs
(thank
you!).  To help us review your code more quickly, **if you are a new
contributor (less than 3 PRs merged) we ask that you complete the
following
tasks and include the filled-out checklist in your PR description.**

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->

# New contributor declaration
- [x] I am not making a trivial change, such as fixing a typo in a
comment.

- [x] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

- Select one of the following.
  - [ ] I have added tests.
    - `/test` for `lit` tests
    - `/unittest` for C++ tests
    - `/python/test` for end-to-end tests
- [x] This PR does not need a test because `test/Triton/combine.mlir
already checks` (NFC change).

- Select one of the following.
  - [x] I have not added any `lit` tests.
- [ ] The `lit` tests I have added follow these [best
practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices),
including the "tests should be minimal" section. (Usually running Python
code
    and using the instructions it generates is not minimal.)

Author: hgreving2304

lib/Dialect/Triton/Transforms/Combine.cpp

@@ -231,6 +231,43 @@ class RankedReduceDescriptorLoads : public mlir::OpRewritePattern<ReshapeOp> {
   }
 };
 
+template <typename OpTy>
+class CombineDotAddPattern : public mlir::OpRewritePattern<OpTy> {
+public:
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(OpTy addOp, mlir::PatternRewriter &rewriter) const override {
+    auto dotOp = addOp.getRhs().template getDefiningOp<DotOp>();
+    bool isDotLHS = false;
+    if (!dotOp) {
+      dotOp = addOp.getLhs().template getDefiningOp<DotOp>();
+      if (!dotOp) {
+        return failure();
+      }
+      isDotLHS = true;
+    }
+    if (!dotOp->hasOneUse()) {
+      return failure();
+    }
+    if (!isZero(dotOp.getC()))
+      return failure();
+    rewriter.modifyOpInPlace(dotOp, [&] {
+      dotOp.getCMutable().assign(isDotLHS ? addOp.getRhs() : addOp.getLhs());
+      dotOp->moveBefore(addOp);
+    });
+    rewriter.replaceAllUsesWith(addOp, dotOp.getResult());
+    return success();
+  }
+};
+
+// AddIOp(DotOp(a, b, c), d) and c==0 => DotOp(a, b, d)
+// AddFOp(DotOp(a, b, c), d) and c==0 => DotOp(a, b, d)
+// AddIOp(d, DotOp(a, b, c)) and c==0 => DotOp(a, b, d)
+// AddFOp(d, DotOp(a, b, c)) and c==0 => DotOp(a, b, d)
+using CombineDotAddIPattern = CombineDotAddPattern<arith::AddIOp>;
+using CombineDotAddFPattern = CombineDotAddPattern<arith::AddFOp>;
+
 } // anonymous namespace
 
 class CombineOpsPass : public impl::TritonCombineOpsBase<CombineOpsPass> {
@@ -240,12 +277,8 @@ class CombineOpsPass : public impl::TritonCombineOpsBase<CombineOpsPass> {
     RewritePatternSet patterns(context);
     ModuleOp m = getOperation();
 
-    // Dot Add %{
     patterns.add<CombineDotAddIPattern>(context);
     patterns.add<CombineDotAddFPattern>(context);
-    patterns.add<CombineDotAddIRevPattern>(context);
-    patterns.add<CombineDotAddFRevPattern>(context);
-    // %}
     patterns.add<CombineSelectMaskedLoadPattern>(context);
     patterns.add<CombineAddPtrPattern>(context);
     patterns.add<CombineBroadcastMulReducePattern>(context);

lib/Dialect/Triton/Transforms/Combine.td

@@ -5,36 +5,6 @@ include "mlir/Dialect/Arith/IR/ArithOps.td"
 include "triton/Dialect/Triton/IR/TritonOps.td"
 include "mlir/IR/PatternBase.td"
 
-
-// AddIOp(DotOp(a, b, c), d) and c==0 => DotOp(a, b, d)
-// AddFOp(DotOp(a, b, c), d) and c==0 => DotOp(a, b, d)
-
-// AddIOp(d, DotOp(a, b, c)) and c==0 => DotOp(a, b, d)
-// AddFOp(d, DotOp(a, b, c)) and c==0 => DotOp(a, b, d)
-def CombineDotAddIPattern : Pat<
-        (Arith_AddIOp $d, (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $overflow),
-        (TT_DotOp $a, $b, $d, $inputPrecision, $maxNumImpreciseAcc, (location $res)),
-        [(Constraint<CPred<"isZero($0)">> $c),
-         (Constraint<CPred<"res->hasOneUse()">, "dot result has a single use">)]>;
-def CombineDotAddFPattern : Pat<
-        (Arith_AddFOp $d, (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $fastmath),
-        (TT_DotOp $a, $b, $d, $inputPrecision, $maxNumImpreciseAcc, (location $res)),
-        [(Constraint<CPred<"isZero($0)">> $c),
-         (Constraint<CPred<"::llvm::cast<::mlir::IntegerAttr>($0).getInt() == 0">> $maxNumImpreciseAcc),
-         (Constraint<CPred<"res->hasOneUse()">, "dot result has a single use">)]>;
-
-def CombineDotAddIRevPattern : Pat<
-        (Arith_AddIOp (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $d, $overflow),
-        (TT_DotOp $a, $b, $d, $inputPrecision, $maxNumImpreciseAcc, (location $res)),
-        [(Constraint<CPred<"isZero($0)">> $c),
-         (Constraint<CPred<"res->hasOneUse()">, "dot result has a single use">)]>;
-def CombineDotAddFRevPattern : Pat<
-        (Arith_AddFOp (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $d, $fastmath),
-        (TT_DotOp $a, $b, $d, $inputPrecision, $maxNumImpreciseAcc, (location $res)),
-        [(Constraint<CPred<"isZero($0)">> $c),
-         (Constraint<CPred<"::llvm::cast<::mlir::IntegerAttr>($0).getInt() == 0">> $maxNumImpreciseAcc),
-         (Constraint<CPred<"res->hasOneUse()">, "dot result has a single use">)]>;
-
 // addptr(addptr(%ptr, %idx0), %idx1) => addptr(%ptr, AddI(%idx0, %idx1))
 //   Note: leave (sub %c0, %c0) canceling to ArithDialect
 //         (ref: ArithCanonicalization.td)