Commit ed8222b2 authored by aartbik's avatar aartbik
Browse files

[mlir] [VectorOps] Implement vector tuple get folding 

Summary: Rewrites get-i tup<a1,...,an> into ai

Reviewers: nicolasvasilache, rriddle, andydavis1

Reviewed By: nicolasvasilache, rriddle, andydavis1

Subscribers: merge_guards_bot, mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, liufengdb, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D73213
parent 4ed7355e

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

+1 −0
Original line number Diff line number Diff line
@@ -1115,6 +1115,7 @@ def Vector_TupleGetOp : 
    }

    static StringRef getIndexAttrName() { return "index"; }

  }];

  let hasFolder = 1;

}



def Vector_PrintOp :

mlir/lib/Dialect/VectorOps/VectorOps.cpp

+12 −0
Original line number Diff line number Diff line
@@ -1681,6 +1681,18 @@ static LogicalResult verify(TupleGetOp op) { 
  return success();

}



OpFoldResult TupleGetOp::fold(ArrayRef<Attribute> operands) {

  // Rewrite:

  //    %t = vector.tuple .., %e_i, ..

  //    %x = vector.tuple_get %t, i

  // into:

  //    %t = vector.tuple .., %e_i, ..  // one less use

  //    %x = %e_i

  if (auto tupleOp = dyn_cast_or_null<TupleOp>(getOperand().getDefiningOp()))

    return tupleOp.getOperand(getIndex());

  return {};

}



//===----------------------------------------------------------------------===//

// ConstantMaskOp

//===----------------------------------------------------------------------===//

mlir/test/Dialect/VectorOps/vector-transforms.mlir

+9 −0
Original line number Diff line number Diff line
@@ -302,3 +302,12 @@ func @vector_transfers(%arg0: index, %arg1: index) { 
  }

  return

}



// CHECK-LABEL: func @tuple_get(%arg0: vector<4xf32>, %arg1: vector<8xf32>)

//       CHECK: return %arg1



func @tuple_get(%arg0: vector<4xf32>, %arg1: vector<8xf32>) -> vector<8xf32> {

  %0 = vector.tuple %arg0, %arg1 : vector<4xf32>, vector<8xf32>

  %1 = vector.tuple_get %0, 1 : tuple<vector<4xf32>, vector<8xf32>>

  return %1 : vector<8xf32>

}