Merge branch 'develop' into dpp-refactor-blockwise-reduce

Author: stefankoncarevic

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

@@ -423,21 +423,6 @@ def MIGraphX_PoolingOp :
   }];
 }
 
-def MIGraphX_FlattenOp :
-    MIGraphX_Op<"flatten">,
-    Arguments<(ins AnyMIXRShaped:$input,
-                   I64Attr:$axis
-                   )>,
-	Results<(outs AnyMIXRShaped:$output)> {
-  let summary = "Flatten tensor";
-  let description = [{
-    The `migraphx.flatten` op.
-  }];
-  let assemblyFormat = [{
-    $input attr-dict `:` type($input) `->` type($output)
-  }];
-}
-
 def MIGraphX_TransposeOp :
     MIGraphX_Op<"transpose">,
     Arguments<(ins AnyMIXRShaped:$input,

mlir/lib/Dialect/Rock/Tuning/RockTuningImpl.cpp

@@ -269,6 +269,8 @@ static void createGemmGemmTuningRangeGreedyPhase1(
   std::mt19937 rng(seed);
   int64_t waveSize =
       rock::lookupArchInfo(rock::getArchValue(gemmGemmOp)).waveSize;
+  bool isWmma = archInfo.isWmma(gemmGemmOp);
+  int64_t numEUPerCU = archInfo.numEUPerCU;
 
   int64_t outputSwizzle{2}, wavesPerEU{0};
   for (uint32_t gemm0MPerBlock : params[0]) {
@@ -302,6 +304,13 @@ static void createGemmGemmTuningRangeGreedyPhase1(
           uint32_t splitKFactor =
               optimalSplitKFactors[rng() % optimalSplitKFactors.size()];
 
+          if (isWmma) {
+            int64_t rdnaWaves = (gemm0MPerBlock / gemmMPerWave) *
+                                (gemm0NPerBlock / gemmNPerWave);
+            if (rdnaWaves < numEUPerCU)
+              continue;
+          }
+
           auto gemmGemmParams = GemmGemmParamsAttr::get(
               gemmGemmOp.getContext(), gemm0MPerBlock, gemm1MPerBlock,
               gemm0NPerBlock, gemmKPerBlock, gemmMPerWave, gemmNPerWave,
@@ -335,6 +344,7 @@ createGemmGemmTuningRangeGreedyPhase2(TuningParamSet *newSpace,
   bool isWmma = archInfo.isWmma(gemmGemmOp);
   int64_t waveSize =
       rock::lookupArchInfo(rock::getArchValue(gemmGemmOp)).waveSize;
+  int64_t numEUPerCU = archInfo.numEUPerCU;
   int64_t outputSwizzle{2}, wavesPerEU{0};
   OpBuilder b(gemmGemmOp.getContext());
 
@@ -356,6 +366,12 @@ createGemmGemmTuningRangeGreedyPhase2(TuningParamSet *newSpace,
   for (uint32_t gemmKPerBlock : validRangeGemmGemmParams[2]) {
     for (uint32_t gemmMPerWave : mPerWaveRange) {
       for (uint32_t gemmNPerWave : nPerWaveRange) {
+        if (isWmma) {
+          int64_t rdnaWaves =
+              (gemm0MPerBlock / gemmMPerWave) * (gemm0NPerBlock / gemmNPerWave);
+          if (rdnaWaves < numEUPerCU)
+            continue;
+        }
         for (uint32_t gemmMnPerXdl : validRangeGemmGemmParams[3]) {
           for (uint32_t gemmKPack : validRangeGemmGemmParams[4]) {
             for (int64_t splitKFactor : optimalSplitKFactors) {

mlir/test/Conversion/MIGraphXToLinalg/migraphx-to-linalg-not-implemented.mlir

@@ -73,12 +73,6 @@ func.func @func_pooling(%arg0: !migraphx.shaped<1x1xf32, 1x1>, %arg1: !migraphx.
   func.return
 }
 
-func.func @func_flatten(%arg0: !migraphx.shaped<1x1xf32, 1x1>, %arg1: !migraphx.shaped<1x1xf32, 1x1>) {
-  // expected-error @+1{{failed to legalize operation 'migraphx.flatten'}}
-  migraphx.flatten %arg0 {axis = 0 : i64}: <1x1xf32, 1x1> -> <1xf32, 1>
-  func.return
-}
-
 func.func @func_slice(%arg0: !migraphx.shaped<1x1xf32, 1x1>, %arg1: !migraphx.shaped<1x1xf32, 1x1>) {
   // expected-error @+1{{failed to legalize operation 'migraphx.slice'}}
   migraphx.slice %arg0 {axes = [0], ends = [1], starts = [0]}: <1x1xf32, 1x1>  -> <1x1xf32, 1x1>