From f0d43a29e30da8f2758e17e1c242917288ab9a0b Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <nicolas.vasilache@gmail.com>
Date: Wed, 23 Jun 2021 09:03:08 +0000
Subject: [PATCH] [mlir][LLVMIR] Fold ExtractValueOp coming from InsertValueOp

Differential Revision: https://reviews.llvm.org/D104769
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td |  1 +
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp  | 10 +++++++++
 mlir/test/Dialect/LLVMIR/canonicalize.mlir  | 23 +++++++++++++++++++++
 3 files changed, 34 insertions(+)
 create mode 100644 mlir/test/Dialect/LLVMIR/canonicalize.mlir

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index e1a32e6d6f15..f714126ce920 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -534,6 +534,7 @@ def LLVM_ExtractValueOp : LLVM_Op<"extractvalue", [NoSideEffect]> {
   let builders = [LLVM_OneResultOpBuilder];
   let parser = [{ return parseExtractValueOp(parser, result); }];
   let printer = [{ printExtractValueOp(p, *this); }];
+  let hasFolder = 1;
 }
 def LLVM_InsertElementOp : LLVM_Op<"insertelement", [NoSideEffect]> {
   let arguments = (ins LLVM_AnyVector:$vector, LLVM_PrimitiveType:$value,
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 512a0ab898e2..95f346096064 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1050,6 +1050,16 @@ static ParseResult parseExtractValueOp(OpAsmParser &parser,
   return success();
 }
 
+OpFoldResult LLVM::ExtractValueOp::fold(ArrayRef<Attribute> operands) {
+  auto insertValueOp = container().getDefiningOp<InsertValueOp>();
+  while (insertValueOp) {
+    if (position() == insertValueOp.position())
+      return insertValueOp.value();
+    insertValueOp = insertValueOp.container().getDefiningOp<InsertValueOp>();
+  }
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // Printing/parsing for LLVM::InsertElementOp.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
new file mode 100644
index 000000000000..a19587318eb2
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt -canonicalize %s -split-input-file | FileCheck %s
+
+// CHECK-LABEL: fold_extractvalue
+llvm.func @fold_extractvalue() -> i32 {
+  //  CHECK-DAG: %[[C0:.*]] = constant 0 : i32
+  %c0 = constant 0 : i32
+  //  CHECK-DAG: %[[C1:.*]] = constant 1 : i32
+  %c1 = constant 1 : i32
+
+  %0 = llvm.mlir.undef : !llvm.struct<(i32, i32)>
+
+  // CHECK-NOT: insertvalue
+  %1 = llvm.insertvalue %c0, %0[0] : !llvm.struct<(i32, i32)>
+  %2 = llvm.insertvalue %c1, %1[1] : !llvm.struct<(i32, i32)>
+
+  // CHECK-NOT: extractvalue
+  %3 = llvm.extractvalue %2[0] : !llvm.struct<(i32, i32)>
+  %4 = llvm.extractvalue %2[1] : !llvm.struct<(i32, i32)>
+
+  //     CHECK: llvm.add %[[C0]], %[[C1]]
+  %5 = llvm.add %3, %4 : i32
+  llvm.return %5 : i32
+}