Compute gridSize properly instead of hard-coding.

llvm · Sep 24, 2020 · a81d05c · a81d05c
1 parent 0bb3f1b
commit a81d05c
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/mlir/lib/Dialect/MIOpen/Transforms/AffixTuningParameters.cpp b/mlir/lib/Dialect/MIOpen/Transforms/AffixTuningParameters.cpp
@@ -164,9 +164,11 @@ void AffixTuningParameters::affixTuningParametersImpl(T &op) {
     validParams.gemmNPerThread = 64;
     validParams.blockSize = 256;
 
-    // XXX. fix gridSize.
-    // need to use (M/MPerBlock)*(N/NPerBlock).
-    gridSize = 784;
+    int64_t M = op.filter().getType().template dyn_cast<MemRefType>().getShape()[1];
+    int64_t N = op.input().getType().template dyn_cast<MemRefType>().getShape()[1];
+
+    // Compute gridSize.
+    gridSize = (M / validParams.gemmMPerBlock) * (N / validParams.gemmNPerBlock);
   }
 
   if (launchDimCallback) {