From da99faf698fe3e7e48a128af5ab37ff378d02876 Mon Sep 17 00:00:00 2001
From: Thomas Raoux <thomas.raoux@openai.com>
Date: Fri, 7 Jun 2024 11:31:08 -0700
Subject: [PATCH] [NFC] Generalize some DotTrait handling (#4098)

Remove some unnecessary restrictions on code related to DotLike.

Only Hopper Dot can take shared memory operands so remove some redundant
checks.
---
 include/triton/Dialect/Triton/IR/Traits.h                | 2 +-
 lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp | 5 -----
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/include/triton/Dialect/Triton/IR/Traits.h b/include/triton/Dialect/Triton/IR/Traits.h
index a768f279446c..6e554aed2358 100644
--- a/include/triton/Dialect/Triton/IR/Traits.h
+++ b/include/triton/Dialect/Triton/IR/Traits.h
@@ -71,7 +71,7 @@ class DotLike : public TraitBase<ConcreteType, DotLike> {
       return op->emitOpError("expected 3 operands");
     auto aTy = cast<TensorOrMemDesc>(op->getOperand(0).getType());
     auto bTy = cast<TensorOrMemDesc>(op->getOperand(1).getType());
-    auto cTy = cast<TensorType>(op->getOperand(2).getType());
+    auto cTy = cast<TensorOrMemDesc>(op->getOperand(2).getType());
     auto aShape = aTy.getShape();
     auto bShape = bTy.getShape();
     auto cShape = cTy.getShape();
diff --git a/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp b/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp
index f41b8ef8687f..6d8279795209 100644
--- a/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp
+++ b/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp
@@ -214,11 +214,6 @@ class FuseTransHopper : public OpRewritePattern<LocalAllocOp> {
       return failure();
 
     auto dot = *allocOp->getUsers().begin();
-    auto dotEnc = dyn_cast<NvidiaMmaEncodingAttr>(
-        cast<RankedTensorType>(dot->getResult(0).getType()).getEncoding());
-    if (!dotEnc || dotEnc.getVersionMajor() != 3)
-      return failure();
-
     if (!allocOp.getSrc())
       return failure();