[mlir][LLVM] Add operand bundle support

Author: Lancern

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

@@ -1210,4 +1210,33 @@ def WorkgroupAttributionAttr
   let assemblyFormat = "`<` $num_elements `,` $element_type `>`";
 }
 
+//===----------------------------------------------------------------------===//
+// OperandBundleAttr
+//===----------------------------------------------------------------------===//
+
+def LLVM_OperandBundleAttr : LLVM_Attr<"OperandBundle", "opbundle"> {
+  let summary = "Operand bundle information";
+  let description = [{
+    Provide information about a single operand bundle. Each operand bundle has a
+    string tag together with various number of SSA value uses. The SSA values
+    are specified through indices into the operation's operand bundle operands.
+  }];
+
+  let parameters = (ins "StringAttr":$tag,
+                        OptionalArrayRefParameter<"uint32_t">:$argIndices);
+  let assemblyFormat = [{
+    `<` $tag (`,` $argIndices^)? `>`
+  }];
+}
+
+def LLVM_OperandBundlesAttr : LLVM_Attr<"OperandBundles", "opbundles"> {
+  let summary = "A list of operand bundle attributes";
+  let description = "A list of operand bundle attributes";
+
+  let parameters = (ins ArrayRefParameter<"OperandBundleAttr">:$bundles);
+  let assemblyFormat = [{
+    `<` $bundles `>`
+  }];
+}
+
 #endif // LLVMIR_ATTRDEFS

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

@@ -550,8 +550,10 @@ def LLVM_InvokeOp : LLVM_Op<"invoke", [
                    Variadic<LLVM_Type>:$callee_operands,
                    Variadic<LLVM_Type>:$normalDestOperands,
                    Variadic<LLVM_Type>:$unwindDestOperands,
+                   Variadic<LLVM_Type>:$bundle_operands,
                    OptionalAttr<DenseI32ArrayAttr>:$branch_weights,
-                   DefaultValuedAttr<CConv, "CConv::C">:$CConv);
+                   DefaultValuedAttr<CConv, "CConv::C">:$CConv,
+                   OptionalAttr<LLVM_OperandBundlesAttr>:$op_bundles);
   let results = (outs Optional<LLVM_Type>:$result);
   let successors = (successor AnySuccessor:$normalDest,
                               AnySuccessor:$unwindDest);
@@ -587,7 +589,8 @@ def LLVM_LandingpadOp : LLVM_Op<"landingpad"> {
 //===----------------------------------------------------------------------===//
 
 def LLVM_CallOp : LLVM_MemAccessOpBase<"call",
-                    [DeclareOpInterfaceMethods<FastmathFlagsInterface>,
+                    [AttrSizedOperandSegments,
+                     DeclareOpInterfaceMethods<FastmathFlagsInterface>,
                      DeclareOpInterfaceMethods<CallOpInterface>,
                      DeclareOpInterfaceMethods<SymbolUserOpInterface>,
                      DeclareOpInterfaceMethods<BranchWeightOpInterface>]> {
@@ -633,6 +636,7 @@ def LLVM_CallOp : LLVM_MemAccessOpBase<"call",
   dag args = (ins OptionalAttr<TypeAttrOf<LLVM_FunctionType>>:$var_callee_type,
                   OptionalAttr<FlatSymbolRefAttr>:$callee,
                   Variadic<LLVM_Type>:$callee_operands,
+                  Variadic<LLVM_Type>:$bundle_operands,
                   DefaultValuedAttr<LLVM_FastmathFlagsAttr,
                                    "{}">:$fastmathFlags,
                   OptionalAttr<DenseI32ArrayAttr>:$branch_weights,
@@ -641,7 +645,8 @@ def LLVM_CallOp : LLVM_MemAccessOpBase<"call",
                   OptionalAttr<LLVM_MemoryEffectsAttr>:$memory_effects,
                   OptionalAttr<UnitAttr>:$convergent,
                   OptionalAttr<UnitAttr>:$no_unwind,
-                  OptionalAttr<UnitAttr>:$will_return
+                  OptionalAttr<UnitAttr>:$will_return,
+                  OptionalAttr<LLVM_OperandBundlesAttr>:$op_bundles
                   );
   // Append the aliasing related attributes defined in LLVM_MemAccessOpBase.
   let arguments = !con(args, aliasAttrs);
@@ -662,6 +667,7 @@ def LLVM_CallOp : LLVM_MemAccessOpBase<"call",
     OpBuilder<(ins "LLVMFunctionType":$calleeType, "StringRef":$callee,
                    CArg<"ValueRange", "{}">:$args)>
   ];
+  let hasVerifier = 1;
   let hasCustomAssemblyFormat = 1;
   let extraClassDeclaration = [{
     /// Returns the callee function type.
@@ -1875,21 +1881,28 @@ def LLVM_InlineAsmOp : LLVM_Op<"inline_asm", [DeclareOpInterfaceMethods<MemoryEf
 
 def LLVM_CallIntrinsicOp
     : LLVM_Op<"call_intrinsic",
-              [DeclareOpInterfaceMethods<FastmathFlagsInterface>]> {
+              [AttrSizedOperandSegments,
+               DeclareOpInterfaceMethods<FastmathFlagsInterface>]> {
   let summary = "Call to an LLVM intrinsic function.";
   let description = [{
     Call the specified llvm intrinsic. If the intrinsic is overloaded, use
     the MLIR function type of this op to determine which intrinsic to call.
   }];
   let arguments = (ins StrAttr:$intrin, Variadic<LLVM_Type>:$args,
+                       Variadic<LLVM_Type>:$bundle_operands,
                        DefaultValuedAttr<LLVM_FastmathFlagsAttr,
-                                         "{}">:$fastmathFlags);
+                                         "{}">:$fastmathFlags,
+                       OptionalAttr<LLVM_OperandBundlesAttr>:$op_bundles);
   let results = (outs Optional<LLVM_Type>:$results);
   let llvmBuilder = [{
     return convertCallLLVMIntrinsicOp(op, builder, moduleTranslation);
   }];
   let assemblyFormat = [{
-    $intrin `(` $args `)` `:` functional-type($args, $results) attr-dict
+    $intrin `(` $args `)`
+    ( `bundlearg` `(` $bundle_operands^ `)` )?
+    `:` functional-type($args, $results)
+    ( `,` `tuple` `<` type($bundle_operands)^ `>` )?
+    attr-dict
   }];
 
   let hasVerifier = 1;

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

@@ -544,6 +544,10 @@ struct CallOpInterfaceLowering : public ConvertOpToLLVMPattern<CallOpType> {
         callOp.getLoc(), packedResult ? TypeRange(packedResult) : TypeRange(),
         promoted, callOp->getAttrs());
 
+    newOp->setAttr(newOp.getOperandSegmentSizesAttrName(),
+                   rewriter.getDenseI32ArrayAttr(
+                       {static_cast<int32_t>(promoted.size()), 0}));
+
     SmallVector<Value, 4> results;
     if (numResults < 2) {
       // If < 2 results, packing did not do anything and we can just return.

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

@@ -837,17 +837,25 @@ class FunctionCallPattern
   matchAndRewrite(spirv::FunctionCallOp callOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (callOp.getNumResults() == 0) {
-      rewriter.replaceOpWithNewOp<LLVM::CallOp>(
+      auto newOp = rewriter.replaceOpWithNewOp<LLVM::CallOp>(
           callOp, std::nullopt, adaptor.getOperands(), callOp->getAttrs());
+      newOp->setAttr(
+          newOp.getOperandSegmentSizesAttrName(),
+          rewriter.getDenseI32ArrayAttr(
+              {static_cast<int32_t>(adaptor.getOperands().size()), 0}));
       return success();
     }
 
     // Function returns a single result.
     auto dstType = typeConverter.convertType(callOp.getType(0));
     if (!dstType)
       return rewriter.notifyMatchFailure(callOp, "type conversion failed");
-    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
+    auto newOp = rewriter.replaceOpWithNewOp<LLVM::CallOp>(
         callOp, dstType, adaptor.getOperands(), callOp->getAttrs());
+    newOp->setAttr(
+        newOp.getOperandSegmentSizesAttrName(),
+        rewriter.getDenseI32ArrayAttr(
+            {static_cast<int32_t>(adaptor.getOperands().size()), 0}));
     return success();
   }
 };