Allow dot operand hoisting for math dialect, arith.truncf, arith.trunci

openxla · Oct 17, 2024 · 8727650 · 8727650
1 parent 1ef30e9
commit 8727650
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp b/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp
@@ -131,19 +131,21 @@ class HoistLayoutConversion : public OpRewritePattern<ConvertLayoutOp> {
     // bitwidth is unable to realize that there is a mixed-precision dot
     // (hence kWidth = 1) but wants to hoist through the type conversion.
     if (isa<arith::ExtFOp>(src) && dotOpEnc.getKWidth() == 1)
-        return failure();
+      return failure();
 
-    // Only consider custom conversions or arith ops.
+    // Only consider custom conversions, math or arith ops.
     // TODO(jlebar): Is this too restrictive?
     if (!isa<FpToFpOp, BitcastOp>(src) && !isPureUnaryInlineAsm(src) &&
-        src->getDialect()->getTypeID() != TypeID::get<arith::ArithDialect>())
+        src->getDialect()->getTypeID() != TypeID::get<arith::ArithDialect>() &&
+        src->getDialect()->getTypeID() != TypeID::get<math::MathDialect>())
       return failure();
 
     // Currently, these instructions are not supported during lowering of
     // shared -> dot_operand layout. Not all types and type conversions are
     // supported.
-    if (isa<arith::TruncIOp, arith::TruncFOp, arith::SelectOp>(src))
+    if (isa<arith::SelectOp>(src)) {
       return failure();
+    }
 
     // Don't hoist through u1 -> fp casts as they aren't supported in
     // ElementwiseOpToLLVM::reorderValues().