[MLIR] [Vector] optimization VectorEmulateMaskedLoadStore crashes

[thaddywu]

To reproduce: https://godbolt.org/z/hKGMY475h
Calling VectorEmulateMaskedLoadStore crashes on this snippet of code
mlir-opt --test-vector-emulate-masked-load-store --allow-unregistered-dialect

"builtin.module"() ({
  "func.func"()<{
    function_type = (memref<12x32xf32>, index, index, vector<32xi1>, vector<32xf32>) -> (),
    sym_name = "fold_vector_maskedstore"
  }>({
  ^bb0(%arg0: memref<12x32xf32>, %arg1: index, %arg2: index, %arg3: vector<32xi1>, %arg4: vector<32xf32>):
    %0 = "memref.subview"(%arg0, %arg1, %arg2)<{
      operandSegmentSizes = array<i32:1,2,0,0>,
      static_offsets = array<i64:-9223372036854775808, -9223372036854775808>,
      static_sizes = array<i64:1, 1>,
      static_strides = array<i64:1, 1>
    }> : (memref<12x32xf32>, index, index) -> memref<f32, strided<[], offset: ?>>
    
    "vector.maskedstore"(%0, %arg3, %arg4) : (
      memref<f32, strided<[], offset: ?>>,
      vector<32xi1>,
      vector<32xf32>
    ) -> ()
    
    "func.return"() : () -> ()
  }) : () -> ()
}) : () -> ()

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Yaoxuan Wu (thaddywu)

Details

To reproduce: https://godbolt.org/z/hKGMY475h Calling VectorEmulateMaskedLoadStore crashes on this snippet of code ` --test-vector-emulate-masked-load-store --allow-unregistered-dialect --mlir-print-ir-before=test-vector-emulate-masked-load-store --mlir-print-ir-after=test-vector-emulate-masked-load-store`

"builtin.module"() ({
  "func.func"()<{
    function_type = (memref<12x32xf32>, index, index, vector<32xi1>, vector<32xf32>) -> (),
    sym_name = "fold_vector_maskedstore"
  }>({
  ^bb0(%arg0: memref<12x32xf32>, %arg1: index, %arg2: index, %arg3: vector<32xi1>, %arg4: vector<32xf32>):
    %0 = "memref.subview"(%arg0, %arg1, %arg2)<{
      operandSegmentSizes = array<i32:1,2,0,0>,
      static_offsets = array<i64:-9223372036854775808, -9223372036854775808>,
      static_sizes = array<i64:1, 1>,
      static_strides = array<i64:1, 1>
    }> : (memref<12x32xf32>, index, index) -> memref<f32, strided<[], offset: ?>>
    
    "vector.maskedstore"(%0, %arg3, %arg4) : (
      memref<f32, strided<[], offset: ?>>,
      vector<32xi1>,
      vector<32xf32>
    ) -> ()
    
    "func.return"() : () -> ()
  }) : () -> ()
}) : () -> ()

[meshtag]

It is doing a masked store on a memref of size 1 with mask and value vector of size 32. This can be potentially UB I think (depending on the mask vector, see the op spec here: https://mlir.llvm.org/docs/Dialects/Vector/#vectormaskedstore-vectormaskedstoreop). Quoting the relevant line from the op spec:

If a mask bit is set and the corresponding index is out-of-bounds for the given base, the behavior is undefined.

The crash can still be fixed though. For example in the above case, the IR would be legal if the mask and value vector size was also 1 (we should probably not crash here).

[meshtag]

Created a PR for the crash here: #173325