Fix vectorization pass for pooling max and update related files

Pooling-Nhwc-max-vec.mlir: examples/BuddyNext PoolingNhwcMaxVectorization.cpp: midend/lib/Conversion/ConvVectorization Pooling-Nhwc-Max-Vectorization.mlir: tests/Conversion
buddy-compiler · Jan 7, 2025 · 35d4de3 · 35d4de3
1 parent 40c48a0
commit 35d4de3
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/examples/BuddyNext/pooling-nhwc-max-vec.mlir b/examples/BuddyNext/pooling-nhwc-max-vec.mlir
@@ -58,9 +58,9 @@ module {
               }
               affine.yield %6 : vector<32xf32>
             }
-           vector.store %5, %arg2[%arg3, %arg4, %arg5, %arg6] : memref<?x?x?x?xf32>, vector<32xf32>
-            %dim_4_next = arith.addi %dim_4, %vl_step : index
-            scf.yield %dim_4_next : index
+            vector.store %5, %arg2[%arg3, %arg4, %arg5, %arg6] : memref<?x?x?x?xf32>, vector<32xf32>
+            %idx_next = arith.addi %arg6, %vl_step : index
+            scf.yield %idx_next : index
           } 
           // Compute the tail size and Process the remaining elements 
           // using masked vector operations.

diff --git a/midend/lib/Conversion/ConvVectorization/PoolingNhwcMaxVectorization.cpp b/midend/lib/Conversion/ConvVectorization/PoolingNhwcMaxVectorization.cpp
@@ -223,7 +223,7 @@ class PoolingNhwcMaxVectorizationPattern : public ConversionPattern {
                     loc, tmp0.getResult(0), output,
                     ValueRange{ivs[0], ivs[1], ivs[2], iv});
                 Value idx =
-                    builder.create<arith::AddIOp>(loc, itrArgs[0], vlStep);
+                    builder.create<arith::AddIOp>(loc, iv, vlStep);
                 builder.create<scf::YieldOp>(loc, idx);
               });
           // Compute the tail size and Process the remaining elements

diff --git a/tests/Conversion/pooling-nhwc-max-vectorization.mlir b/tests/Conversion/pooling-nhwc-max-vectorization.mlir
@@ -20,7 +20,7 @@
 // CHECK-NEXT:           affine.yield %12 : vector<16xf32>
 // CHECK-NEXT:         }
 // CHECK-NEXT:         vector.store %9, %arg2[%arg3, %arg4, %arg5, %arg6] : memref<?x?x?x?xf32>, vector<16xf32>
-// CHECK-NEXT:         %10 = arith.addi %arg7, %c16 : index
+// CHECK-NEXT:         %10 = arith.addi %arg6, %c16 : index
 // CHECK-NEXT:         scf.yield %10 : index
 // CHECK-NEXT:       }