Merge c9e273d into a6c5ebf

iree-org · Jun 17, 2024 · 321942b · 321942b
2 parents a6c5ebf + c9e273d
commit 321942b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/compiler/src/iree/compiler/GlobalOptimization/SetEncoding.cpp b/compiler/src/iree/compiler/GlobalOptimization/SetEncoding.cpp
@@ -128,6 +128,16 @@ static MatmulNarrowSizes getMatmulNarrowSizes(ShapedType outType,
   if (!ShapedType::isDynamic(N) && N < kNarrowThreshold) {
     narrow.N = N;
   }
+
+  // Only pick 1 if both are present
+  if (narrow.M && narrow.N) {
+    if (*narrow.M <= *narrow.N) {
+      narrow.N.reset();
+    } else {
+      narrow.M.reset();
+    }
+  }
+
   return narrow;
 }
 

diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/set_encoding.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/set_encoding.mlir
@@ -1021,9 +1021,9 @@ util.func public @batch_matmul_f32f32f32_narrow_MN(%arg0 : tensor<64x4x250xf32>,
 //  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 //      CHECK: util.func public @batch_matmul_f32f32f32_narrow_MN(
-//      CHECK:   iree_encoding.upper_bound_tile_size tensor<64x4x250xf32, #iree_encoding.encoding<role =  LHS, element_types = [f32, f32, f32], matmul_narrow_M = 4 : index, matmul_narrow_N = 2 : index, user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]>>
-//      CHECK:   iree_encoding.upper_bound_tile_size tensor<64x250x2xf32, #iree_encoding.encoding<role =  RHS, element_types = [f32, f32, f32], matmul_narrow_M = 4 : index, matmul_narrow_N = 2 : index, user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]>>
-//      CHECK:   iree_encoding.upper_bound_tile_size tensor<64x4x2xf32, #iree_encoding.encoding<role =  RESULT, element_types = [f32, f32, f32], matmul_narrow_M = 4 : index, matmul_narrow_N = 2 : index, user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]>>
+//      CHECK:   iree_encoding.upper_bound_tile_size tensor<64x4x250xf32, #iree_encoding.encoding<role =  LHS, element_types = [f32, f32, f32], matmul_narrow_N = 2 : index, user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]>>
+//      CHECK:   iree_encoding.upper_bound_tile_size tensor<64x250x2xf32, #iree_encoding.encoding<role =  RHS, element_types = [f32, f32, f32], matmul_narrow_N = 2 : index, user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]>>
+//      CHECK:   iree_encoding.upper_bound_tile_size tensor<64x4x2xf32, #iree_encoding.encoding<role =  RESULT, element_types = [f32, f32, f32], matmul_narrow_N = 2 : index, user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]>>
 //      CHECK:   linalg.batch_matmul
 
 // -----