[Mlir-commits] [mlir] f0d43a2 - [mlir][LLVMIR] Fold ExtractValueOp coming from InsertValueOp

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-06-23T10:04:24Z
New Revision: f0d43a29e30da8f2758e17e1c242917288ab9a0b

URL: https://github.com/llvm/llvm-project/commit/f0d43a29e30da8f2758e17e1c242917288ab9a0b
DIFF: https://github.com/llvm/llvm-project/commit/f0d43a29e30da8f2758e17e1c242917288ab9a0b.diff

LOG: [mlir][LLVMIR] Fold ExtractValueOp coming from InsertValueOp

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

Added: 
    mlir/test/Dialect/LLVMIR/canonicalize.mlir

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
    mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index e1a32e6d6f153..f714126ce9209 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 512a0ab898e2d..95f3460960649 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 0000000000000..a19587318eb22
--- /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
+}