[mlir][Vector] Fold extractelement splat.

[jacquesguan]

Summary

This revision supports to fold vector.extractelement (splat X) -> X.

[ftynse]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

954–956

Can't we? If we are extracting the element from a splat, isn't it always the same regardless of the position?

958

Nit: please use full sentences in comments, that is, start with a capital letter and end with a dot.

[jacquesguan]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

954–956

But if the position exceeds the range of the vector, the result should be a poison value? Could we ignore this possibility?

[ftynse]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

954–956

Good question. There is no poison semantics in MLIR. I would interpret accessing out of bounds as C-style UB, at which point folding it to the splat value is as valid as any other behavior.

[ThomasRaoux]

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

954–956

I agree. Note that llvm also does this optimization in this case:
https://github.com/llvm/llvm-project/blob/17fdaccccfad9b143e4aadbcdda7f645de127153/llvm/lib/Analysis/InstructionSimplify.cpp#L4779

[jacquesguan]

Address comment.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

954–956

Done, thanks.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 
 OpFoldResult vector::ExtractElementOp::fold(ArrayRef<Attribute> operands) {
   // Skip the 0-D vector here now.
   if (operands.size() < 2)
     return {};
 
   Attribute src = operands[0];
   Attribute pos = operands[1];
-  if (!src || !pos)
+
+  // We could not fold if the position is not constant.
+  if (!pos)
+    return {};
+
+  // fold extractelement (splat X) -> X
+  if (auto splat = getVector().getDefiningOp<vector::SplatOp>())
+    return splat.getInput();
+
+  if (!src)
     return {};
 
   auto srcElements = src.cast<DenseElementsAttr>().getValues<Attribute>();
 
   auto attr = pos.dyn_cast<IntegerAttr>();
   uint64_t posIdx = attr.getInt();
 
   return srcElements[posIdx];

mlir/test/Dialect/Vector/canonicalize.mlir

 //       CHECK:   %[[C:.+]] = arith.constant 5 : i32
 //       CHECK:   return %[[C]]
 func @extract_element_fold() -> i32 {
   %v = arith.constant dense<[1, 3, 5, 7]> : vector<4xi32>
   %i = arith.constant 2 : i32
   %1 = vector.extractelement %v[%i : i32] : vector<4xi32>
   return %1 : i32
 }
+
+// CHECK-LABEL: func @extract_element_splat_fold
+//  CHECK-SAME: (%[[ARG:.+]]: i32)
+//       CHECK:   return %[[ARG]]
+func @extract_element_splat_fold(%a : i32) -> i32 {
+  %v = vector.splat %a : vector<4xi32>
+  %i = arith.constant 2 : i32
+  %1 = vector.extractelement %v[%i : i32] : vector<4xi32>
+  return %1 : i32
+}