Additional vectorization logic for Input tensor (Matrix B).

llvm · Jun 6, 2020 · ff3096e · ff3096e
1 parent 587540d
commit ff3096e
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/mlir/lib/Dialect/MIOpenOps/CppOutput/gridwise_convolution_implicit_gemm_v4r4_gen_xdlops.cpp b/mlir/lib/Dialect/MIOpenOps/CppOutput/gridwise_convolution_implicit_gemm_v4r4_gen_xdlops.cpp
@@ -177,6 +177,16 @@ class TunableParameters : public TunableParametersBase {
         vectorizableLength = ctx.c;
         gemmKVectorizable = true;
         // gemmN dimension non-vectorizable.
+      } else if (ctx.dimCI == 0) {
+        if(ctx.y == 1 && ctx.x == 1 && ctx.strideH == 1 && ctx.strideW == 1 && ctx.paddingHL == 0 &&
+           ctx.paddingHR == 0 && ctx.paddingWL == 0 && ctx.paddingWR == 0) {
+            // \todo there are more configs that can go through this if branch
+            srcDataPerRead_Gemm = gcd(srcDataPerRead_Gemm, ctx.n * ctx.hi * ctx.wi);
+
+            gemmNVectorizable = true;
+        } else {
+            srcDataPerRead_Gemm = 1;
+        }
       } else if (ctx.dimHI == 2 && ctx.dimWI == 3) {
         if(ctx.y == 1 && ctx.x == 1 && ctx.strideH == 1 && ctx.strideW == 1 && ctx.paddingHL == 0 &&
            ctx.paddingHR == 0 && ctx.paddingWL == 0 && ctx.paddingWR == 0) {