Skip to content

Commit 178ebd4

Browse files
authored
[ElementwiseOpToLLVM] Do not call 'maybeDeduplicate' for dpas encoding (#3508)
Closes #3380 The problem occurs when we try to lower an `arith.andi` op in a chain of `tl.dot + tl.where + tl.dot`. The `maybeDeduplicate` function seems to handle dot/dpas layouts incorrectly. When trying to deduplicate mask values for each 16-col block it seems to mishandle strides & constancy somehow leading to the mask values for the second 16-col block to be ignored. Signed-off-by: dchigarev <[email protected]>
1 parent eb7cf23 commit 178ebd4

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#ifndef TRITON_CONVERSION_TRITONGPU_TO_ELEMENTWISE_OP_H
22
#define TRITON_CONVERSION_TRITONGPU_TO_ELEMENTWISE_OP_H
33

4+
#include "intel/include/Dialect/TritonIntelGPU/IR/Attributes.h"
45
#include "mlir/Conversion/LLVMCommon/Pattern.h"
56
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
67
#include "mlir/Support/LLVM.h"
@@ -86,7 +87,8 @@ class ElementwiseOpConversionBase : public ConvertOpToLLVMPattern<SourceOp> {
8687
return resultVals;
8788
while (auto sliced = dyn_cast<SliceEncodingAttr>(baseEncoding))
8889
baseEncoding = sliced.getParent();
89-
if (isa<LinearEncodingAttr, DotOperandEncodingAttr>(baseEncoding)) {
90+
if (isa<LinearEncodingAttr, DotOperandEncodingAttr,
91+
intel::DpasEncodingAttr>(baseEncoding)) {
9092
// TODO: this logic seems incorrect for mma layout. Skip for now.
9193
// The following test crashes and some other miscompile:
9294
// test_core::test_fp8_dot_acc

0 commit comments

Comments
 (0)