diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp index 7b976943e1bb..704308329861 100644 --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -908,8 +908,9 @@ mlir::affine::hasNoInterveningEffect &opsToErase, PostDominanceInfo &postDominanceInfo, @@ -936,6 +937,16 @@ static void findUnusedStore(AffineWriteOpInterface writeA, if (srcAccess != destAccess) continue; + // Check that the store types match. If types differ, writeB may not cover + // all bytes written by writeA (e.g. a narrower vector type), so + // conservatively assume writeA is not dead. + // One could be tempted whether writeA type is smaller than writeB, however + // it can become tricky with cases like vector<4xi6> vs vector<3xi8> due to + // padding that can be datalayout dependent. + if (writeA.getValueToStore().getType() != + writeB.getValueToStore().getType()) + continue; + // writeB must postdominate writeA. if (!postDominanceInfo.postDominates(writeB, writeA)) continue; diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir index 901e15911f4e..fb6eef941790 100644 --- a/mlir/test/Dialect/Affine/scalrep.mlir +++ b/mlir/test/Dialect/Affine/scalrep.mlir @@ -997,3 +997,40 @@ func.func @zero_d_memrefs() { } return } + +// CHECK-LABEL: func @vector_store_dead_elim_different_types +// A vector_store with a different vector type at the same base index must NOT +// cause the earlier vector_store to be treated as a dead store. +// (GitHub issue #113687) +func.func @vector_store_dead_elim_different_types(%arg0: memref<20x1xi64>) { + %c0 = arith.constant 0 : index + // CHECK-DAG: %[[CST1:.+]] = arith.constant dense<1> + // CHECK-DAG: %[[CST2:.+]] = arith.constant dense<2> + %cst1 = arith.constant dense<1> : vector<5xi64> + %cst2 = arith.constant dense<2> : vector<1xi64> + // Both stores must be preserved: the second (narrow) store does not fully + // overwrite the first (wide) store. + // CHECK: affine.vector_store %[[CST1]], %arg0[%c0, %c0] : memref<20x1xi64>, vector<5xi64> + affine.vector_store %cst1, %arg0[%c0, %c0] : memref<20x1xi64>, vector<5xi64> + // CHECK: affine.vector_store %[[CST2]], %arg0[%c0, %c0] : memref<20x1xi64>, vector<1xi64> + affine.vector_store %cst2, %arg0[%c0, %c0] : memref<20x1xi64>, vector<1xi64> + return +} + +// CHECK-LABEL: func @vector_store_dead_elim_same_type +// A vector_store with the same type at the same base index should still cause +// the earlier store to be treated as a dead store. +func.func @vector_store_dead_elim_same_type(%arg0: memref<20x1xi64>) { + %c0 = arith.constant 0 : index + %cst1 = arith.constant dense<1> : vector<5xi64> + %cst2 = arith.constant dense<2> : vector<5xi64> + // CHECK-DAG: %[[CST2:.+]] = arith.constant dense<2> + // The first store is dead — the second store (same type, same index) + // fully overwrites it. + // CHECK-NOT: arith.constant dense<1> + // CHECK-NOT: affine.vector_store {{.*}} dense<1> + // CHECK: affine.vector_store %[[CST2]], %arg0[%c0, %c0] : memref<20x1xi64>, vector<5xi64> + affine.vector_store %cst1, %arg0[%c0, %c0] : memref<20x1xi64>, vector<5xi64> + affine.vector_store %cst2, %arg0[%c0, %c0] : memref<20x1xi64>, vector<5xi64> + return +}