Commit 12b49518 authored by Andrzej Warzynski's avatar Andrzej Warzynski
Browse files

[mlir][vector] Add missing support for scalable vectors 

This patch adds the missing logic so that the
`TransferReadPermutationLowering` can be used for scalable vectors. To
this end:
  * TransferOp custom C++ builder is updated to support scalable
    vectors,
  * `TransferOpReduceRank` is also updated to support scalable vectors.

This pattern is relevant when lowering `linalg.matmul` via
`vector_multi_reduction` for scalable vectors.

I've also updated relevant code in `TransferOpReduceRank` not to use
`llvm::to_vector` for constructing `SmallVector` from `ArrayRef`. That
hook doesn't work for `ArraryRef<bool>` (*), so for consistency I
switched to an explicit constructor (so that both `newShape` and
`newScalableDim` are constructed in a similar fashion).

(*) IIUC, that's due how implicit narrowing conversions between `bool`
and `*bool` work. Note that these narrowing conversions change when
using initializer lists, see
  * https://en.cppreference.com/w/cpp/language/list_initialization.

Depends on D157092

Differential Revision: https://reviews.llvm.org/D157268
parent f8087884

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

+6 −2
Original line number Diff line number Diff line
@@ -4949,11 +4949,15 @@ void vector::TransposeOp::build(OpBuilder &builder, OperationState &result, 
                                Value vector, ArrayRef<int64_t> transp) {

  VectorType vt = llvm::cast<VectorType>(vector.getType());

  SmallVector<int64_t, 4> transposedShape(vt.getRank());

  for (unsigned i = 0; i < transp.size(); ++i)

  SmallVector<bool, 4> transposedScalableDims(vt.getRank());

  for (unsigned i = 0; i < transp.size(); ++i) {

    transposedShape[i] = vt.getShape()[transp[i]];

    transposedScalableDims[i] = vt.getScalableDims()[transp[i]];

  }



  result.addOperands(vector);

  result.addTypes(VectorType::get(transposedShape, vt.getElementType()));

  result.addTypes(VectorType::get(transposedShape, vt.getElementType(),

                                  transposedScalableDims));

  result.addAttribute(TransposeOp::getTranspAttrName(result.name),

                      builder.getI64ArrayAttr(transp));

}

mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp

+10 −5
Original line number Diff line number Diff line
@@ -115,8 +115,11 @@ struct TransferReadPermutationLowering 
    // Apply the reverse transpose to deduce the type of the transfer_read.

    ArrayRef<int64_t> originalShape = op.getVectorType().getShape();

    SmallVector<int64_t> newVectorShape(originalShape.size());

    ArrayRef<bool> originalScalableDims = op.getVectorType().getScalableDims();

    SmallVector<bool> newScalableDims(originalShape.size());

    for (const auto &pos : llvm::enumerate(permutation)) {

      newVectorShape[pos.value()] = originalShape[pos.index()];

      newScalableDims[pos.value()] = originalScalableDims[pos.index()];

    }



    // Transpose in_bounds attribute.
@@ -126,8 +129,8 @@ struct TransferReadPermutationLowering 
                         : ArrayAttr();



    // Generate new transfer_read operation.

    VectorType newReadType =

        VectorType::get(newVectorShape, op.getVectorType().getElementType());

    VectorType newReadType = VectorType::get(

        newVectorShape, op.getVectorType().getElementType(), newScalableDims);

    Value newRead = rewriter.create<vector::TransferReadOp>(

        op.getLoc(), newReadType, op.getSource(), op.getIndices(),

        AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
@@ -345,14 +348,16 @@ struct TransferOpReduceRank : public OpRewritePattern<vector::TransferReadOp> { 
      return success();

    }



    SmallVector<int64_t> newShape = llvm::to_vector<4>(

    SmallVector<int64_t> newShape(

        originalVecType.getShape().take_back(reducedShapeRank));

    SmallVector<bool> newScalableDims(

        originalVecType.getScalableDims().take_back(reducedShapeRank));

    // Vector rank cannot be zero. Handled by TransferReadToVectorLoadLowering.

    if (newShape.empty())

      return rewriter.notifyMatchFailure(op, "rank-reduced vector is 0-d");



    VectorType newReadType =

        VectorType::get(newShape, originalVecType.getElementType());

    VectorType newReadType = VectorType::get(

        newShape, originalVecType.getElementType(), newScalableDims);

    ArrayAttr newInBoundsAttr =

        op.getInBounds()

            ? rewriter.getArrayAttr(

mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir

+26 −2
Original line number Diff line number Diff line 
// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s



// CHECK-LABEL: func @lower_permutation_with_mask(

// CHECK-LABEL: func @lower_permutation_with_mask_fixed_width(

//       CHECK:   %[[vec:.*]] = arith.constant dense<-2.000000e+00> : vector<7x1xf32>

//       CHECK:   %[[mask:.*]] = arith.constant dense<[true, false, true, false, true, true, true]> : vector<7xi1>

//       CHECK:   %[[b:.*]] = vector.broadcast %[[mask]] : vector<7xi1> to vector<1x7xi1>

//       CHECK:   %[[tp:.*]] = vector.transpose %[[b]], [1, 0] : vector<1x7xi1> to vector<7x1xi1>

//       CHECK:   vector.transfer_write %[[vec]], %{{.*}}[%{{.*}}, %{{.*}}], %[[tp]] {in_bounds = [false, true]} : vector<7x1xf32>, memref<?x?xf32>

func.func @lower_permutation_with_mask(%A : memref<?x?xf32>, %base1 : index,

func.func @lower_permutation_with_mask_fixed_width(%A : memref<?x?xf32>, %base1 : index,

                                       %base2 : index) {

  %fn1 = arith.constant -2.0 : f32

  %vf0 = vector.splat %fn1 : vector<7xf32>
@@ -17,6 +17,30 @@ func.func @lower_permutation_with_mask(%A : memref<?x?xf32>, %base1 : index, 
  return

}



// CHECK-LABEL:   func.func @permutation_with_mask_scalable(

// CHECK-SAME:      %[[ARG_0:.*]]: memref<?x?xf32>,

// CHECK-SAME:      %[[IDX_1:.*]]: index,

// CHECK-SAME:      %[[IDX_2:.*]]: index) -> vector<8x[4]x2xf32> {

// CHECK:           %[[C0:.*]] = arith.constant 0 : index

// CHECK:           %[[PASS_THROUGH:.*]] = arith.constant 0.000000e+00 : f32

// CHECK:           %[[MASK:.*]] = vector.create_mask %[[IDX_2]], %[[IDX_1]] : vector<2x[4]xi1>

// CHECK:           %[[T_READ:.*]] = vector.transfer_read %[[ARG_0]]{{\[}}%[[C0]], %[[C0]]], %[[PASS_THROUGH]], %[[MASK]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x[4]xf32>

// CHECK:           %[[BCAST:.*]] = vector.broadcast %[[T_READ]] : vector<2x[4]xf32> to vector<8x2x[4]xf32>

// CHECK:           %[[TRANSPOSE:.*]] = vector.transpose %[[BCAST]], [0, 2, 1] : vector<8x2x[4]xf32> to vector<8x[4]x2xf32>

// CHECK:           return %[[TRANSPOSE]] : vector<8x[4]x2xf32>

// CHECK:         }

func.func @permutation_with_mask_scalable(%2: memref<?x?xf32>, %dim_1: index, %dim_2: index) -> (vector<8x[4]x2xf32>) {



  %c0 = arith.constant 0 : index

  %cst_0 = arith.constant 0.000000e+00 : f32



  %mask = vector.create_mask %dim_2, %dim_1 : vector<2x[4]xi1>

  %1 = vector.transfer_read %2[%c0, %c0], %cst_0, %mask 

    {in_bounds = [true, true, true], permutation_map = affine_map<(d0, d1) -> (0, d1, d0)>}

    : memref<?x?xf32>, vector<8x[4]x2xf32>

  return %1 : vector<8x[4]x2xf32>

}



transform.sequence failures(propagate) {

^bb1(%module_op: !transform.any_op):

  %f = transform.structured.match ops{["func.func"]} in %module_op