Merge commit 'a6d11f755d6c01c6ef369af5d4b4f3447522313b'

Author: anmyachev

.gitignore

@@ -35,7 +35,9 @@ python/triton/backends/*
 !python/triton/backends/driver.py
 
 # Language extras
-python/triton/language/extra
+python/triton/language/extra/*
+!python/triton/language/extra/__init__.py
+!python/triton/language/extra/libdevice.py
 
 # Tools extras
 python/triton/tools/extra

include/triton/Dialect/TritonGPU/Transforms/DecomposeScaledBlocked.h

@@ -1,7 +1,40 @@
 #include "mlir/IR/PatternMatch.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
 
 namespace mlir::triton::gpu {
 
+class DecomposeScaledBlocked : public OpRewritePattern<DotScaledOp> {
+public:
+  DecomposeScaledBlocked(MLIRContext *context, PatternBenefit benefit)
+      : OpRewritePattern<DotScaledOp>(context, benefit) {}
+
+  LogicalResult matchAndRewrite(DotScaledOp scaledDotOp,
+                                PatternRewriter &rewriter) const override;
+
+protected:
+  FloatType getComputeType(ScaleDotElemType aType, ScaleDotElemType bType,
+                           PatternRewriter &rewriter) const;
+  TypedValue<RankedTensorType> scaleTo16(PatternRewriter &rewriter,
+                                         TypedValue<RankedTensorType> scale,
+                                         FloatType computeType) const;
+  TypedValue<RankedTensorType>
+  broadcastScale(PatternRewriter &rewriter, DotScaledOp scaledDotOp,
+                 ModuleOp mod, TypedValue<RankedTensorType> scale,
+                 int dim) const;
+  TypedValue<RankedTensorType> maskNan(PatternRewriter &rewriter,
+                                       DotScaledOp scaledDotOp, ModuleOp mod,
+                                       TypedValue<RankedTensorType> mxfp,
+                                       TypedValue<RankedTensorType> scale,
+                                       int dim) const;
+  TypedValue<RankedTensorType> scaleArg(PatternRewriter &rewriter,
+                                        DotScaledOp scaledDotOp, int opIdx,
+                                        FloatType computeType) const;
+  TypedValue<RankedTensorType>
+  cvtDotOperand(PatternRewriter &rewriter, DotScaledOp scaledDotOp, int opIdx,
+                TypedValue<RankedTensorType> v) const;
+  static SmallVector<int, 2> getTransposeOrder(int rank);
+};
+
 void populateDecomposeScaledBlockedPatterns(mlir::RewritePatternSet &patterns,
                                             int benefit);