[mlir][LLVMIR] Allow insert/extract element ops to have a constant index

[Mogball]

Summary

This patch allows LLVM's insertelement and extractelement ops to
optionally have a constant index instead of a dynamic index, allowing
syntax like

%0 = llvm.extractelement %v[1] : vector<2xi32>
%1 = llvm.insertelement %0, %v[2] : vector<2xi32>

Depends on D131971

[rriddle]

This feels like it's missing justification, can you add why? What's the benefit/necessity of the complexity here?

[Mogball]

LLVMIR allows these ops to have inline constant indices. I thought the need to materialize an extra constant in these cases was a little verbose.

[rriddle]

In D131981#3732380, @Mogball wrote:

LLVMIR allows these ops to have inline constant indices. I thought the need to materialize an extra constant in these cases was a little verbose.

I'm not sure that design principle scales though, given that LLVM technically allows inline constants on everything.

[Mogball]

Yeah now that I think about it, you're right. LLVMIR constants are different than MLIR constants.

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

 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
 def LLVM_ExtractElementOp : LLVM_Op<"extractelement", [NoSideEffect,
     TypesMatchWith<"result type matches vector element type", "vector", "res",
                    "LLVM::getVectorElementType($_self)">]> {
   let summary = "Extract an element from an LLVM vector.";
 
-  let arguments = (ins LLVM_AnyVector:$vector, AnyInteger:$position);
+  let arguments = (ins LLVM_AnyVector:$vector,
+                       Optional<AnyInteger>:$dynamicPosition,
+                       OptionalAttr<IndexAttr>:$constPosition);
   let results = (outs LLVM_Type:$res);
 
   let builders = [
     OpBuilder<(ins "Value":$vector, "Value":$position,
-                   CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
+                   CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
+    OpBuilder<(ins "Type":$res, "Value":$vector, "Value":$position)>,
+    OpBuilder<(ins "Value":$vector, "int64_t":$position)>,
   ];
 
   let assemblyFormat = [{
-    $vector `[` $position `:` type($position) `]` attr-dict `:` type($vector)
+    $vector `[`
+      ($dynamicPosition^ `:` type($dynamicPosition)) : (`` $constPosition) ?
+    `]` attr-dict `:` type($vector)
   }];
 
+  let hasVerifier = 1;
+
   string llvmBuilder = [{
-    $res = builder.CreateExtractElement($vector, $position);
+    if (op.getDynamicPosition())
+      $res = builder.CreateExtractElement($vector, $dynamicPosition.front());
+    else
+      $res = builder.CreateExtractElement(
+          $vector, builder.getInt($constPosition.value()));
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractValueOp
 //===----------------------------------------------------------------------===//
 
 def LLVM_ExtractValueOp : LLVM_Op<"extractvalue", [NoSideEffect]> {

 
 def LLVM_InsertElementOp : LLVM_Op<"insertelement", [NoSideEffect,
     TypesMatchWith<"argument type matches vector element type", "vector",
                    "value", "LLVM::getVectorElementType($_self)">,
     AllTypesMatch<["res", "vector"]>]> {
   let summary = "Insert an element into an LLVM vector.";
 
   let arguments = (ins LLVM_AnyVector:$vector, LLVM_PrimitiveType:$value,
-                       AnyInteger:$position);
+                       Optional<AnyInteger>:$dynamicPosition,
+                       OptionalAttr<IndexAttr>:$constPosition);
   let results = (outs LLVM_AnyVector:$res);
 
-  let builders = [LLVM_OneResultOpBuilder];
+  let builders = [
+    LLVM_OneResultOpBuilder,
+    OpBuilder<(ins "Type":$res, "Value":$vector, "Value":$value,
+                   "Value":$position)>,
+    OpBuilder<(ins "Value":$vector, "Value":$value, "Value":$position)>,
+    OpBuilder<(ins "Value":$vector, "Value":$value, "int64_t":$position)>,
+  ];
 
   let assemblyFormat = [{
-    $value `,` $vector `[` $position `:` type($position) `]` attr-dict `:`
-    type($vector)
+    $value `,` $vector `[`
+      ($dynamicPosition^ `:` type($dynamicPosition)) : (`` $constPosition) ?
+    `]` attr-dict `:` type($vector)
   }];
 
+  let hasVerifier = 1;
+
   string llvmBuilder = [{
-    $res = builder.CreateInsertElement($vector, $value, $position);
+    if (op.getDynamicPosition())
+      $res = builder.CreateInsertElement($vector, $value,
+                                         $dynamicPosition.front());
+    else
+      $res = builder.CreateInsertElement(
+          $vector, $value, builder.getInt(op.getConstPosition().value()));
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // InsertValueOp
 //===----------------------------------------------------------------------===//
 
 def LLVM_InsertValueOp : LLVM_Op<

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

ParseResult CallOp::parse(OpAsmParser &parser, OperationState &result) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
-/// Expects vector to be an LLVM vector type and position to be an integer type.
-void LLVM::ExtractElementOp::build(OpBuilder &b, OperationState &result,
-                                   Value vector, Value position,
-                                   ArrayRef<NamedAttribute> attrs) {
-  auto vectorType = vector.getType();
-  auto llvmType = LLVM::getVectorElementType(vectorType);
-  build(b, result, llvmType, vector, position);
-  result.addAttributes(attrs);
+/// Verify the dynamic and constant indices of an insert or extract element
+/// operation.
+static LogicalResult verifyElementIndex(Operation *op, Value dynamicPosition,
+                                        IntegerAttr constPosition,
+                                        Value vector) {
+  if (dynamicPosition && constPosition)
+    return op->emitOpError(
+        "expected either a dynamic position or constant position, not both");
+  else if (!dynamicPosition && !constPosition)
+    return op->emitOpError(
+        "expected a dynamic position or a constant position");
+
+  if (constPosition) {
+    int64_t pos = constPosition.getInt();
+    if (pos < 0 || pos >= vector.getType().cast<VectorType>().getNumElements())
+      return op->emitOpError("constant index out-of-range");
+  }
+  return success();
+}
+
+LogicalResult ExtractElementOp::verify() {
+  return verifyElementIndex(*this, getDynamicPosition(), getConstPositionAttr(),
+                            getVector());
+}
+
+void ExtractElementOp::build(OpBuilder &builder, OperationState &state,
+  build(builder, state, LLVM::getVectorElementType(vector.getType()), vector,
+        position);
+  state.addAttributes(attrs);
+}
+
+void ExtractElementOp::build(OpBuilder &builder, OperationState &state,
+                             Type res, Value vector, Value position) {
+  build(builder, state, res, vector, /*dynamicPosition=*/position,
+        /*constPosition=*/nullptr);
+}
+
+void ExtractElementOp::build(OpBuilder &builder, OperationState &state,
+                             Value vector, int64_t position) {
+  build(builder, state, LLVM::getVectorElementType(vector.getType()), vector,
+        /*dynamicPosition=*/nullptr,
+        /*constPosition=*/builder.getIndexAttr(position));
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractValueOp
 //===----------------------------------------------------------------------===//
 
 /// Extract the type at `position` in the LLVM IR aggregate type
 /// `containerType`. Each element of `position` is an index into a nested

 void ExtractValueOp::build(OpBuilder &builder, OperationState &state,
                            Value container, ArrayRef<int64_t> position) {
   build(builder, state,
         getInsertExtractValueElementType(container.getType(), position),
         container, builder.getAttr<DenseI64ArrayAttr>(position));
 }
 
 //===----------------------------------------------------------------------===//
+// InsertElementOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult InsertElementOp::verify() {
+  return verifyElementIndex(*this, getDynamicPosition(), getConstPositionAttr(),
+                            getVector());
+}
+
+void InsertElementOp::build(OpBuilder &builder, OperationState &state, Type res,
+                            Value vector, Value value, Value position) {
+  build(builder, state, res, vector, value, /*dynamicPosition=*/position,
+        /*constPosition=*/nullptr);
+}
+
+void InsertElementOp::build(OpBuilder &builder, OperationState &state,
+                            Value vector, Value value, Value position) {
+  build(builder, state, vector.getType(), vector, value, position);
+}
+
+void InsertElementOp::build(OpBuilder &builder, OperationState &state,
+                            Value vector, Value value, int64_t position) {
+  build(builder, state, vector.getType(), vector, value,
+        /*dynamicPosition=*/nullptr,
+        /*constPosition=*/builder.getIndexAttr(position));
+}
+
+//===----------------------------------------------------------------------===//
 // InsertValueOp
 //===----------------------------------------------------------------------===//
 
 /// Infer the value type from the container type and position.
 static ParseResult
 parseInsertExtractValueElementType(AsmParser &parser, Type &valueType,
                                    Type containerType,
                                    DenseI64ArrayAttr position) {

mlir/test/Dialect/LLVMIR/invalid.mlir

 func.func @invalid_vector_type_5(%a : vector<4xf32>, %idx : i32) -> vector<4xf32> {
   // expected-error@+1 {{failed to verify that argument type matches vector element type}}
   %b = "llvm.insertelement"(%a, %a, %idx) : (vector<4xf32>, vector<4xf32>, i32) -> vector<4xf32>
   return %b : vector<4xf32>
 }
 
 // -----
 
+func.func @invalid_element_index(%a : vector<4xf32>, %idx : i32) -> f32 {
+  // expected-error@+1 {{expected either a dynamic position or constant position, not both}}
+  %b = "llvm.extractelement"(%a, %idx) {constPosition = 1 : index} : (vector<4xf32>, i32) -> f32
+  return %b : f32
+}
+
+// -----
+
+func.func @invalid_element_index(%a : vector<4xf32>, %idx : i32) -> f32 {
+  // expected-error@+1 {{expected a dynamic position or a constant position}}
+  %b = "llvm.extractelement"(%a) : (vector<4xf32>) -> f32
+  return %b : f32
+}
+
+// -----
+
+func.func @invalid_element_index(%a : vector<4xf32>, %idx : i32) -> f32 {
+  // expected-error@+1 {{constant index out-of-range}}
+  %b = llvm.extractelement %a[5] : vector<4xf32>
+  return %b : f32
+}
+
+// -----
+
 func.func @null_non_llvm_type() {
   // expected-error@+1 {{custom op 'llvm.mlir.null' invalid kind of type specified}}
   llvm.mlir.null : i32
 }
 
 // -----
 
 func.func @nvvm_invalid_shfl_pred_1(%arg0 : i32, %arg1 : i32, %arg2 : i32, %arg3 : i32) {

mlir/test/Dialect/LLVMIR/roundtrip.mlir

llvm.func @vararg_func(%arg0: i32, ...) {
   llvm.intr.vacopy %3 to %5
   // CHECK: llvm.intr.vaend %[[CAST1]]
   // CHECK: llvm.intr.vaend %[[CAST0]]
   llvm.intr.vaend %5
   llvm.intr.vaend %3
   // CHECK: llvm.return
   llvm.return
 }
+
+// CHECK-LABEL: @insert_extract_element
+// CHECK-SAME: %[[V0:.*]]: vector<4xf32>, %[[V1:.*]]: vector<4xf32>, %[[I:.*]]: i32
+llvm.func @insert_extract_element(%v0: vector<4xf32>, %v1: vector<4xf32>, %i: i32) {
+  // CHECK: %[[E0:.*]] = llvm.extractelement %[[V0]][%[[I]] : i32]
+  // CHECK: %[[E1:.*]] = llvm.extractelement %[[V0]][2]
+  // CHECK: %[[E2:.*]] = llvm.insertelement %[[E0]], %[[V1]][%[[I]] : i32]
+  // CHECK: %[[E3:.*]] = llvm.insertelement %[[E1]], %[[E2]][3]
+  %0 = llvm.extractelement %v0[%i : i32] : vector<4xf32>
+  %1 = llvm.extractelement %v0[2] : vector<4xf32>
+  %2 = llvm.insertelement %0, %v1[%i : i32] : vector<4xf32>
+  %3 = llvm.insertelement %1, %2[3] : vector<4xf32>
+  llvm.return
+}