From 2a6234d614c0bf4a4cea7be46c3e1372f2616374 Mon Sep 17 00:00:00 2001
From: Victor Perez <victor.perez@codeplay.com>
Date: Wed, 30 Oct 2024 09:47:42 +0000
Subject: [PATCH] [OptRed] Do not assert on reductions after 3D dot ops

Triton supports 3D dot operations and #2518 enables them. Avoid
assertions when this pass finds reductions of 3D tensors.

Signed-off-by: Victor Perez <victor.perez@codeplay.com>
---
 .../TritonIntelGPUTransforms/OptimizeReductionLocality.cpp    | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/OptimizeReductionLocality.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/OptimizeReductionLocality.cpp
index ccc564ea66..3bc0b7f0ab 100644
--- a/third_party/intel/lib/TritonIntelGPUTransforms/OptimizeReductionLocality.cpp
+++ b/third_party/intel/lib/TritonIntelGPUTransforms/OptimizeReductionLocality.cpp
@@ -153,13 +153,15 @@ struct DpasOperandPattern final : OpRewritePattern<ReduceOp> {
     // Check this is has `triton_intel_gpu.dpas` encoding.
     Value operand = operands.front();
     auto type = cast<RankedTensorType>(operand.getType());
+    // Only support reduction after 2D-dot for now.
+    if (type.getRank() != 2)
+      return failure();
     auto encoding =
         llvm::dyn_cast_or_null<DpasEncodingAttr>(type.getEncoding());
     if (!encoding)
       return failure();
 
     // Axis 1 will lead to within-warp reduction.
-    assert(type.getRank() == 2 && "Expecting 2D tensor");
     if (op.getAxis() != preferredReductionAxis)
       return failure();