[mlir][Vector] Prevent AVX2 lowering for non-f32 transpose ops

[dcaballe]

Summary

The AVX2 lowering for transpose operations is only applicable to f32 vector types.

[aartbik]

Just curious if F32 data movement would not just work for any 32-bit entity, but change looks good to me.

Make sure to get presubmit green first before committing though.

[dcaballe]

Thanks!

Just curious if F32 data movement would not just work for any 32-bit entity, but change looks good to me.

The existing patterns currently use f32 specific instructions (note the Ps):

Value t0 = mm256UnpackLoPs(ib, vs[0], vs[1]);
Value t1 = mm256UnpackHiPs(ib, vs[0], vs[1]);
Value t2 = mm256UnpackLoPs(ib, vs[2], vs[3]);
Value t3 = mm256UnpackHiPs(ib, vs[2], vs[3]);
Value t4 = mm256UnpackLoPs(ib, vs[4], vs[5]);
Value t5 = mm256UnpackHiPs(ib, vs[4], vs[5]);
Value t6 = mm256UnpackLoPs(ib, vs[6], vs[7]);
Value t7 = mm256UnpackHiPs(ib, vs[6], vs[7]);

It should be relatively easy to enable them on i32 since they have the same size.

mlir/lib/Dialect/X86Vector/Transforms/AVXTranspose.cpp

TransposeOpLowering(LoweringOptions loweringOptions, MLIRContext *context,
       : OpRewritePattern<vector::TransposeOp>(context, benefit),
         loweringOptions(loweringOptions) {}
 
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     // Check if the source vector type is supported. AVX2 patterns can only be
-    // applied if the vector type has two dimensions greater than one.
+    // applied to f32 vector types with two dimensions greater than one.
     VectorType srcType = op.getVectorType();
+    if (!srcType.getElementType().isF32())
+      return rewriter.notifyMatchFailure(op, "Unsupported vector element type");
+
     SmallVector<int64_t> srcGtOneDims;
     for (auto &en : llvm::enumerate(srcType.getShape()))
       if (en.value() > 1)
         srcGtOneDims.push_back(en.index());
 
     if (srcGtOneDims.size() != 2)
       return rewriter.notifyMatchFailure(op, "Unsupported vector type");
 

mlir/test/Dialect/Vector/vector-transpose-lowering.mlir

func @transpose210_1x8x8xf32(%arg0: vector<1x8x8xf32>) -> vector<8x8x1xf32> {
   // AVX2-NEXT: vector.insert {{.*}}[7]
   // AVX2-NEXT: vector.shape_cast %{{.*}} : vector<8x8xf32> to vector<8x8x1xf32>
   %0 = vector.transpose %arg0, [2, 1, 0] : vector<1x8x8xf32> to vector<8x8x1xf32>
   return %0 : vector<8x8x1xf32>
 }
 
 // -----
 
+func @do_not_lower_nonf32_to_avx2(%arg0: vector<4x8xi32>) -> vector<8x4xi32> {
+  %0 = vector.transpose %arg0, [1, 0] : vector<4x8xi32> to vector<8x4xi32>
+  return %0 : vector<8x4xi32>
+}
+
+// AVX2-NOT: vector.shuffle
+
+// -----
+
 // AVX2-LABEL: func @transpose021_8x1x8
 func @transpose021_8x1x8xf32(%arg0: vector<8x1x8xf32>) -> vector<8x8x1xf32> {
   %0 = vector.transpose %arg0, [0, 2, 1] : vector<8x1x8xf32> to vector<8x8x1xf32>
   return %0 : vector<8x8x1xf32>
 }
 
 // AVX2-NOT: vector.shuffle