[mlir][LLVMIR] Fold ExtractValueOp coming from InsertValueOp

[nicolasvasilache]

Summary

ftynse

[ftynse]

mlir/test/Dialect/LLVMIR/canonicalize.mlir

18

Maybe check the actual operand values?

[nicolasvasilache]

Improve test.

[nicolasvasilache]

Fix test.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_ExtractValueOp : LLVM_Op<"extractvalue", [NoSideEffect]> {
   let arguments = (ins LLVM_AnyAggregate:$container, ArrayAttr:$position);
   let results = (outs LLVM_Type:$res);
   string llvmBuilder = [{
     $res = builder.CreateExtractValue($container, extractPosition($position));
   }];
   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,
                    AnyInteger:$position);
   let results = (outs LLVM_AnyVector:$res);
   string llvmBuilder = [{
     $res = builder.CreateInsertElement($vector, $value, $position);
   }];

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

auto elementType = getInsertExtractValueElementType(
       parser, containerType, positionAttr, attributeLoc, trailingTypeLoc);
   if (!elementType)
     return failure();
 
   result.addTypes(elementType);
   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.
 //===----------------------------------------------------------------------===//
 
 static void printInsertElementOp(OpAsmPrinter &p, InsertElementOp &op) {
   p << op.getOperationName() << ' ' << op.value() << ", " << op.vector() << "["
     << op.position() << " : " << op.position().getType() << "]";
   p.printOptionalAttrDict(op->getAttrs());

mlir/test/Dialect/LLVMIR/canonicalize.mlir

+// RUN: mlir-opt -canonicalize %s -split-input-file | FileCheck %s
+
+// CHECK-LABEL: fold_extractvalue
+llvm.func @fold_extractvalue() -> i32 {
+  %c0 = constant 0 : 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
+  %5 = llvm.add %3, %4 : i32
+  llvm.return %5 : i32
+}