Skip to content

Commit e438919

Browse files
authored
[OptRed] Do not assert on reductions after 3D dot ops (#2590)
Triton supports 3D dot operations and #2518 enables using them with `DpasEncodingAttr` encoding. Avoid assertions when this pass finds reductions of 3D tensors. Signed-off-by: Victor Perez <[email protected]>
1 parent efce869 commit e438919

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

third_party/intel/lib/TritonIntelGPUTransforms/OptimizeReductionLocality.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,13 +153,15 @@ struct DpasOperandPattern final : OpRewritePattern<ReduceOp> {
153153
// Check this is has `triton_intel_gpu.dpas` encoding.
154154
Value operand = operands.front();
155155
auto type = cast<RankedTensorType>(operand.getType());
156+
// Only support reduction after 2D-dot for now.
157+
if (type.getRank() != 2)
158+
return failure();
156159
auto encoding =
157160
llvm::dyn_cast_or_null<DpasEncodingAttr>(type.getEncoding());
158161
if (!encoding)
159162
return failure();
160163

161164
// Axis 1 will lead to within-warp reduction.
162-
assert(type.getRank() == 2 && "Expecting 2D tensor");
163165
if (op.getAxis() != preferredReductionAxis)
164166
return failure();
165167

0 commit comments

Comments
 (0)