diff --git a/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h b/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h index 87c5585b1632..a5f65e77a6d4 100644 --- a/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h +++ b/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h @@ -87,10 +87,12 @@ class ElementwiseOpConversionBase : public ConvertOpToLLVMPattern { if (!encoding) // encoding not available return resultVals; - if (!dyn_cast(encoding) && - !dyn_cast(encoding)) { - // TODO: constraining the ecndoing type here is necessary for avoiding - // crashes in the getElemsPerThread call below happening in the + Attribute baseEncoding = encoding; + while (auto sliced = dyn_cast(baseEncoding)) + baseEncoding = sliced.getParent(); + if (isa(baseEncoding)) { + // TODO: this logic seems incorrect for mma layout. Skip for now. + // The following test crashes and some other miscompile: // test_core::test_fp8_dot_acc return resultVals; }