[mlir][Vector] Fold ShuffleOp if result is identical to one of source vectors.

[jacquesguan]

Summary

For example, we could do the following eliminations:

fold vector.shuffle V1, V2, [0, 1, 2, 3] : <4xi32>, <2xi32> -> V1
fold vector.shuffle V1, V2, [4, 5] : <4xi32>, <2xi32> -> V2

[jacquesguan]

remove some irrelevant change.

[nicolasvasilache]

Thanks!

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

1773

Note: we use snakeCase for function and variable names in he MLIR coding style:

static bool isStepIndexArray(ArrayAttr idxArr, int64_t begin, int64_t width) {...}

1774

int64_t expected = begin;
return idxArr.size() == width && 
   llvm::all_of(idxArr.getAsValueRange<IntegerAttr>(), 
                      [&expected](IntegerAttr attr){ return attr.getZExtValue() == expected++; });

[jacquesguan]

Address comment.

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

1773

Done, thanks a lot.

1774

Done.

[hpmorgan]

Hi @jacquesguan ,

This change has caused a build error when building with clang:

/usr/bin/clang++ -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -D_DEBUG -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -I/mnt/vss/_work/1/b/llvm/Release/tools/mlir/lib/Dialect/Vector/IR -I/mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Vector/IR -I/mnt/vss/_work/1/b/llvm/Release/include -I/mnt/vss/_work/1/llvm-project/llvm/include -I/mnt/vss/_work/1/llvm-project/mlir/include -I/mnt/vss/_work/1/b/llvm/Release/tools/mlir/include -fPIC -fvisibility-inlines-hidden -Werror -Werror=date-time -Werror=unguarded-availability-new -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wmissing-field-initializers -pedantic -Wno-long-long -Wc++98-compat-extra-semi -Wimplicit-fallthrough -Wcovered-switch-default -Wno-noexcept-type -Wnon-virtual-dtor -Wdelete-non-virtual-dtor -Wsuggest-override -Wstring-conversion -Wmisleading-indentation -fdiagnostics-color -ffunction-sections -fdata-sections -Werror=mismatched-tags -Werror=global-constructors -O3 -DNDEBUG  -fno-exceptions -fno-rtti -UNDEBUG -std=c++14 -MD -MT tools/mlir/lib/Dialect/Vector/IR/CMakeFiles/obj.MLIRVector.dir/VectorOps.cpp.o -MF tools/mlir/lib/Dialect/Vector/IR/CMakeFiles/obj.MLIRVector.dir/VectorOps.cpp.o.d -o tools/mlir/lib/Dialect/Vector/IR/CMakeFiles/obj.MLIRVector.dir/VectorOps.cpp.o -c /mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
/mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Vector/IR/VectorOps.cpp:1793:24: error: comparison of integers of different signs: 'size_t' (aka 'unsigned long') and 'int64_t' (aka 'long') [-Werror,-Wsign-compare]
  return idxArr.size() == width &&
         ~~~~~~~~~~~~~ ^  ~~~~~
/mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Vector/IR/VectorOps.cpp:1796:52: error: comparison of integers of different signs: 'uint64_t' (aka 'unsigned long') and 'int64_t' (aka 'long') [-Werror,-Wsign-compare]
                        return attr.getZExtValue() == expected++;
                               ~~~~~~~~~~~~~~~~~~~ ^  ~~~~~~~~~~

Thanks,

Henry

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

ShuffleOp::inferReturnTypes(MLIRContext *, Optional<Location>,
   shape.reserve(v1Type.getRank());
   shape.push_back(std::max<size_t>(1, op.getMask().size()));
   llvm::append_range(shape, v1Type.getShape().drop_front());
   inferredReturnTypes.push_back(
       VectorType::get(shape, v1Type.getElementType()));
   return success();
 }
 
+static bool isStepIndexArray(ArrayAttr idxArr, int64_t begin, int64_t width) {
+  int64_t expected = begin;
+  return idxArr.size() == width &&
+         llvm::all_of(idxArr.getAsValueRange<IntegerAttr>(),
+                      [&expected](auto attr) {
+                        return attr.getZExtValue() == expected++;
+                      });
+}
+
 OpFoldResult vector::ShuffleOp::fold(ArrayRef<Attribute> operands) {
+  // fold shuffle V1, V2, [0, 1, 2, 3] : <4xi32>, <2xi32> -> V1
+  if (!getV1VectorType().isScalable() &&
+      isStepIndexArray(getMask(), 0, getV1VectorType().getDimSize(0)))
+    return getV1();
+  // fold shuffle V1, V2, [4, 5] : <4xi32>, <2xi32> -> V2
+  if (!getV1VectorType().isScalable() && !getV2VectorType().isScalable() &&
+      isStepIndexArray(getMask(), getV1VectorType().getDimSize(0),
+                       getV2VectorType().getDimSize(0)))
+    return getV2();
+
   Attribute lhs = operands.front(), rhs = operands.back();
   if (!lhs || !rhs)
     return {};
 
   auto lhsType = lhs.getType().cast<VectorType>();
   // Only support 1-D for now to avoid complicated n-D DenseElementsAttr
   // manipulation.
   if (lhsType.getRank() != 1)

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 
 func @shuffle_1D_direct(%arg0: vector<2xf32>, %arg1: vector<2xf32>) -> vector<2xf32> {
   %1 = vector.shuffle %arg0, %arg1 [0, 1] : vector<2xf32>, vector<2xf32>
   return %1 : vector<2xf32>
 }
 // CHECK-LABEL: @shuffle_1D_direct(
 // CHECK-SAME: %[[A:.*]]: vector<2xf32>,
 // CHECK-SAME: %[[B:.*]]: vector<2xf32>)
-//       CHECK:   %[[s:.*]] = llvm.shufflevector %[[A]], %[[B]] [0, 1] : vector<2xf32>, vector<2xf32>
-//       CHECK:   return %[[s]] : vector<2xf32>
+//       CHECK:   return %[[A:.*]]: vector<2xf32>
 
 // -----
 
 func @shuffle_1D_index_direct(%arg0: vector<2xindex>, %arg1: vector<2xindex>) -> vector<2xindex> {
   %1 = vector.shuffle %arg0, %arg1 [0, 1] : vector<2xindex>, vector<2xindex>
   return %1 : vector<2xindex>
 }
 // CHECK-LABEL: @shuffle_1D_index_direct(
 // CHECK-SAME: %[[A:.*]]: vector<2xindex>,
 // CHECK-SAME: %[[B:.*]]: vector<2xindex>)
-//       CHECK-DAG:   %[[T0:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<2xindex> to vector<2xi64>
-//       CHECK-DAG:   %[[T1:.*]] = builtin.unrealized_conversion_cast %[[B]] : vector<2xindex> to vector<2xi64>
-//       CHECK:   %[[T2:.*]] = llvm.shufflevector %[[T0]], %[[T1]] [0, 1] : vector<2xi64>, vector<2xi64>
-//       CHECK:   %[[T3:.*]] = builtin.unrealized_conversion_cast %[[T2]] : vector<2xi64> to vector<2xindex>
-//       CHECK:   return %[[T3]] : vector<2xindex>
+//       CHECK:   return  %[[A:.*]]: vector<2xindex>
 
 // -----
 
 func @shuffle_1D(%arg0: vector<2xf32>, %arg1: vector<3xf32>) -> vector<5xf32> {
   %1 = vector.shuffle %arg0, %arg1 [4, 3, 2, 1, 0] : vector<2xf32>, vector<3xf32>
   return %1 : vector<5xf32>
 }
 // CHECK-LABEL: @shuffle_1D(

mlir/test/Dialect/Vector/canonicalize.mlir

 //       CHECK:   %[[V:.+]] = arith.constant dense<[3, 2, 5, 1]> : vector<4xi32>
 //       CHECK:   return %[[V]]
 func @shuffle_1d() -> vector<4xi32> {
   %v0 = arith.constant dense<[0, 1, 2]> : vector<3xi32>
   %v1 = arith.constant dense<[3, 4, 5]> : vector<3xi32>
   %shuffle = vector.shuffle %v0, %v1 [3, 2, 5, 1] : vector<3xi32>, vector<3xi32>
   return %shuffle : vector<4xi32>
 }
+
+// CHECK-LABEL: func @shuffle_fold1
+//       CHECK:   %arg0 : vector<4xi32>
+func @shuffle_fold1(%v0 : vector<4xi32>, %v1 : vector<2xi32>) -> vector<4xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3] : vector<4xi32>, vector<2xi32>
+  return %shuffle : vector<4xi32>
+}
+
+// CHECK-LABEL: func @shuffle_fold2
+//       CHECK:   %arg1 : vector<2xi32>
+func @shuffle_fold2(%v0 : vector<4xi32>, %v1 : vector<2xi32>) -> vector<2xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [4, 5] : vector<4xi32>, vector<2xi32>
+  return %shuffle : vector<2xi32>
+}
+
+// CHECK-LABEL: func @shuffle_fold3
+//       CHECK:   return %arg0 : vector<4x5x6xi32>
+func @shuffle_fold3(%v0 : vector<4x5x6xi32>, %v1 : vector<2x5x6xi32>) -> vector<4x5x6xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3] : vector<4x5x6xi32>, vector<2x5x6xi32>
+  return %shuffle : vector<4x5x6xi32>
+}
+
+// CHECK-LABEL: func @shuffle_fold4
+//       CHECK:   return %arg1 : vector<2x5x6xi32>
+func @shuffle_fold4(%v0 : vector<4x5x6xi32>, %v1 : vector<2x5x6xi32>) -> vector<2x5x6xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [4, 5] : vector<4x5x6xi32>, vector<2x5x6xi32>
+  return %shuffle : vector<2x5x6xi32>
+}
+
+// CHECK-LABEL: func @shuffle_nofold1
+//       CHECK:   %[[V:.+]] = vector.shuffle %arg0, %arg1 [0, 1, 2, 3, 4] : vector<4xi32>, vector<2xi32>
+//       CHECK:   return %[[V]]
+func @shuffle_nofold1(%v0 : vector<4xi32>, %v1 : vector<2xi32>) -> vector<5xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3, 4] : vector<4xi32>, vector<2xi32>
+  return %shuffle : vector<5xi32>
+}
+
+// CHECK-LABEL: func @shuffle_nofold2
+//       CHECK:   %[[V:.+]] = vector.shuffle %arg0, %arg1 [0, 1, 2, 3] : vector<[4]xi32>, vector<[2]xi32>
+//       CHECK:   return %[[V]]
+func @shuffle_nofold2(%v0 : vector<[4]xi32>, %v1 : vector<[2]xi32>) -> vector<4xi32> {
+  %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3] : vector<[4]xi32>, vector<[2]xi32>
+  return %shuffle : vector<4xi32>
+}