From 35d4de398354a5c49d1fa8b94f967772eb6f7925 Mon Sep 17 00:00:00 2001 From: FloatingcloudKnight <1348185166@qq.com> Date: Tue, 7 Jan 2025 11:58:45 +0000 Subject: [PATCH] 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 --- examples/BuddyNext/pooling-nhwc-max-vec.mlir | 6 +++--- .../ConvVectorization/PoolingNhwcMaxVectorization.cpp | 2 +- tests/Conversion/pooling-nhwc-max-vectorization.mlir | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/BuddyNext/pooling-nhwc-max-vec.mlir b/examples/BuddyNext/pooling-nhwc-max-vec.mlir index d6d8a35d14..5cfe2085db 100644 --- 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, 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, 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 index 280d11b226..e07e3fa718 100644 --- 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(loc, itrArgs[0], vlStep); + builder.create(loc, iv, vlStep); builder.create(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 index 0fa188eeb6..919d2626c7 100644 --- 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, 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: }