Merge commit '85c7e15cc2b81f42a6baea850cc802cd47fb5a9c'

intel · Jun 8, 2024 · 7a38e94 · 7a38e94
2 parents e2d00da + 85c7e15
commit 7a38e94
Show file tree

Hide file tree

Showing 4 changed files with 407 additions and 7 deletions.
diff --git 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
@@ -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();