[mlir][LLVMIR] Add support for translating insertelement/extractelement.

Add support for translating llvm::InsertElement and llvm::ExtractElement.

Differential Revision: https://reviews.llvm.org/D125674

2 files changed, 26 insertions, 3 deletions

diff --git a/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp b/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp
index 124d82b3f78c..cc10d0bccfd2 100644
--- a/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp
+++ b/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp
@@ -616,8 +616,7 @@ static StringRef lookupOperationNameFromOpcode(unsigned opcode) {
       INST(Freeze, Freeze), INST(Call, Call),
       // FIXME: select
       // FIXME: vaarg
-      // FIXME: extractelement
-      // FIXME: insertelement
+      INST(ExtractElement, ExtractElement), INST(InsertElement, InsertElement),
       // ShuffleVector is handled specially.
       // InsertValue is handled specially.
       // ExtractValue is handled specially.
@@ -775,7 +774,9 @@ LogicalResult Importer::processInstruction(llvm::Instruction *inst) {
   case llvm::Instruction::IntToPtr:
   case llvm::Instruction::AddrSpaceCast:
   case llvm::Instruction::Freeze:
-  case llvm::Instruction::BitCast: {
+  case llvm::Instruction::BitCast:
+  case llvm::Instruction::ExtractElement:
+  case llvm::Instruction::InsertElement: {
     OperationState state(loc, lookupOperationNameFromOpcode(inst->getOpcode()));
     SmallVector<Value, 4> ops;
     ops.reserve(inst->getNumOperands());
diff --git a/mlir/test/Target/LLVMIR/Import/basic.ll b/mlir/test/Target/LLVMIR/Import/basic.ll
index 24810ac7732b..05c09cebcaa3 100644
--- a/mlir/test/Target/LLVMIR/Import/basic.ll
+++ b/mlir/test/Target/LLVMIR/Import/basic.ll
@@ -583,3 +583,25 @@ define <4 x half> @shuffle_vec(<4 x half>* %arg0, <4 x half>* %arg1) {
   %shuffle = shufflevector <4 x half> %val0, <4 x half> %val1, <4 x i32> <i32 2, i32 3, i32 undef, i32 undef>
   ret <4 x half> %shuffle
 }
+
+; ExtractElement
+; CHECK-LABEL: llvm.func @extract_element
+define half @extract_element(<4 x half>* %vec, i32 %idx) {
+  ; CHECK: %[[V0:.+]] = llvm.load %{{.+}} : !llvm.ptr<vector<4xf16>>
+  %val0 = load <4 x half>, <4 x half>* %vec
+  ; CHECK: %[[V1:.+]] = llvm.extractelement %[[V0]][%{{.+}} : i32] : vector<4xf16>
+  %r = extractelement <4 x half> %val0, i32 %idx
+  ; CHECK: llvm.return %[[V1]]
+  ret half %r
+}
+
+; InsertElement
+; CHECK-LABEL: llvm.func @insert_element
+define <4 x half> @insert_element(<4 x half>* %vec, half %v, i32 %idx) {
+  ; CHECK: %[[V0:.+]] = llvm.load %{{.+}} : !llvm.ptr<vector<4xf16>>
+  %val0 = load <4 x half>, <4 x half>* %vec
+  ; CHECK: %[[V1:.+]] = llvm.insertelement %{{.+}}, %[[V0]][%{{.+}} : i32] : vector<4xf16>
+  %r = insertelement <4 x half> %val0, half %v, i32 %idx
+  ; CHECK: llvm.return %[[V1]]
+  ret <4 x half> %r
+}