[mlir][llvm] Add support for importing masked intrinsics from LLVM IR.

[gysit]

Summary

The revision adds support for importing the masked load/store and
gather/scatter intrinsics from LLVM IR. To enable the import, the
revision also includes an extension of the mlirBuilder code generation
to support variadic arguments.

Depends on D136057

[ftynse]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

247

Why does this have to materialize the vector?

mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp

75

Be aware that isVariableLength checks for both variadic and optional.

[gysit]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

247

True I can return an ArrayRef.

mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp

75

I think we want to allow both at the moment but this is indeed a bit confusing given the name of the method. There are certain operations that use Optional<..> and certain others use Variadic<...>, which makes sense since the ReturnOp either takes 0 or 1 argument.

Would you prefer to support only Variadic operands here and handle Optional<...> in a different way. I can try how many things break if I am more restrictive.

[gysit]

Address comments.

[ftynse]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

247

At this point, do you really need this helper? :)

mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp

75

Yes, I prefer with start with being more strict. Optional operands are more frequent than variadic ones and we would better be on the safe side given we are out of the tested path here.

[gysit]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

247

I'll drop it yes.

[gysit]

Address comments.

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

 def LLVM_GetActiveLaneMaskOp
     : LLVM_OneResultIntrOp<"get.active.lane.mask", [0], [0], [Pure]> {
   let arguments = (ins LLVM_Type:$base, LLVM_Type:$n);
   let assemblyFormat = "$base `,` $n attr-dict `:` "
     "type($base) `,` type($n) `to` type($res)";
 }
 
 /// Create a call to Masked Load intrinsic.
-def LLVM_MaskedLoadOp : LLVM_Op<"intr.masked.load"> {
+def LLVM_MaskedLoadOp : LLVM_OneResultIntrOp<"masked.load"> {
   let arguments = (ins LLVM_Type:$data, LLVM_Type:$mask,
                    Variadic<LLVM_Type>:$pass_thru, I32Attr:$alignment);
   let results = (outs LLVM_AnyVector:$res);
+  let assemblyFormat =
+    "operands attr-dict `:` functional-type(operands, results)";
+
   string llvmBuilder = [{
     $res = $pass_thru.empty() ? builder.CreateMaskedLoad(
         $_resultType, $data, llvm::Align($alignment), $mask) :
       builder.CreateMaskedLoad(
         $_resultType, $data, llvm::Align($alignment), $mask, $pass_thru[0]);
   }];
-  let assemblyFormat =
-    "operands attr-dict `:` functional-type(operands, results)";
+  string mlirBuilder = [{
+    $res = $_builder.create<LLVM::MaskedLoadOp>($_location,
+      $_resultType, $data, $mask, $pass_thru, $_int_attr($alignment));
+  }];
+  list<int> llvmArgIndices = [0, 2, 3, 1];
 }
 
 /// Create a call to Masked Store intrinsic.
-def LLVM_MaskedStoreOp : LLVM_Op<"intr.masked.store"> {
+def LLVM_MaskedStoreOp : LLVM_ZeroResultIntrOp<"masked.store"> {
   let arguments = (ins LLVM_Type:$value, LLVM_Type:$data, LLVM_Type:$mask,
                    I32Attr:$alignment);
   let builders = [LLVM_VoidResultTypeOpBuilder, LLVM_ZeroResultOpBuilder];
+  let assemblyFormat = "$value `,` $data `,` $mask attr-dict `:` "
+    "type($value) `,` type($mask) `into` type($data)";
+
   string llvmBuilder = [{
     builder.CreateMaskedStore(
       $value, $data, llvm::Align($alignment), $mask);
   }];
-  let assemblyFormat = "$value `,` $data `,` $mask attr-dict `:` "
-    "type($value) `,` type($mask) `into` type($data)";
+  string mlirBuilder = [{
+    $_builder.create<LLVM::MaskedStoreOp>($_location,
+      $value, $data, $mask, $_int_attr($alignment));
+  }];
+  list<int> llvmArgIndices = [0, 1, 3, 2];
 }
 
 /// Create a call to Masked Gather intrinsic.
-def LLVM_masked_gather : LLVM_Op<"intr.masked.gather"> {
+def LLVM_masked_gather : LLVM_OneResultIntrOp<"masked.gather"> {
   let arguments = (ins LLVM_AnyVector:$ptrs, LLVM_Type:$mask,
                    Variadic<LLVM_Type>:$pass_thru, I32Attr:$alignment);
   let results = (outs LLVM_Type:$res);
   let builders = [LLVM_OneResultOpBuilder];
+  let assemblyFormat =
+    "operands attr-dict `:` functional-type(operands, results)";
+
   string llvmBuilder = [{
     $res = $pass_thru.empty() ? builder.CreateMaskedGather(
         $_resultType, $ptrs, llvm::Align($alignment), $mask) :
       builder.CreateMaskedGather(
         $_resultType, $ptrs, llvm::Align($alignment), $mask, $pass_thru[0]);
   }];
-  let assemblyFormat =
-    "operands attr-dict `:` functional-type(operands, results)";
+  string mlirBuilder = [{
+    $res = $_builder.create<LLVM::masked_gather>($_location,
+      $_resultType, $ptrs, $mask, $pass_thru, $_int_attr($alignment));
+  }];
+  list<int> llvmArgIndices = [0, 2, 3, 1];
 }
 
 /// Create a call to Masked Scatter intrinsic.
-def LLVM_masked_scatter : LLVM_Op<"intr.masked.scatter"> {
+def LLVM_masked_scatter : LLVM_ZeroResultIntrOp<"masked.scatter"> {
   let arguments = (ins LLVM_Type:$value, LLVM_Type:$ptrs, LLVM_Type:$mask,
                    I32Attr:$alignment);
   let builders = [LLVM_VoidResultTypeOpBuilder, LLVM_ZeroResultOpBuilder];
+  let assemblyFormat = "$value `,` $ptrs `,` $mask attr-dict `:` "
+    "type($value) `,` type($mask) `into` type($ptrs)";
+
   string llvmBuilder = [{
     builder.CreateMaskedScatter(
       $value, $ptrs, llvm::Align($alignment), $mask);
   }];
-  let assemblyFormat = "$value `,` $ptrs `,` $mask attr-dict `:` "
-    "type($value) `,` type($mask) `into` type($ptrs)";
+  string mlirBuilder = [{
+    $_builder.create<LLVM::masked_scatter>($_location,
+      $value, $ptrs, $mask, $_int_attr($alignment));
+  }];
+  list<int> llvmArgIndices = [0, 1, 3, 2];
 }
 
 /// Create a call to Masked Expand Load intrinsic.
 def LLVM_masked_expandload : LLVM_IntrOp<"masked.expandload", [0], [], [], 1> {
   let arguments = (ins LLVM_Type, LLVM_Type, LLVM_Type);
 }
 
 /// Create a call to Masked Compress Store intrinsic.

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

class LLVM_OpBase<Dialect dialect, string mnemonic, list<Trait> traits = []> :
   //             name matches the result name, by a reference to store the
   //             result of the newly created MLIR operation to;
   //   - $_int_attr - substituted by a call to an integer attribute matcher;
   //   - $_resultType - substituted with the MLIR result type;
   //   - $_location - substituted with the MLIR location;
   //   - $_builder - substituted with the MLIR builder;
   //   - $_qualCppClassName - substitiuted with the MLIR operation class name.
   // Additionally, `$$` can be used to produce the dollar character.
-  // FIXME: The $name variable resolution does not support variadic arguments.
   string mlirBuilder = "";
 
   // An array that specifies a mapping from MLIR argument indices to LLVM IR
   // operand indices. The mapping is necessary since argument and operand
   // indices do not always match. If not defined, the array is set to the
   // identity permutation. An operation may define any custom index permutation
   // and set a specific argument index to -1 if it does not map to an LLVM IR
   // operand.

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

 
 /// Converts an array of unsigned indices to a signed integer position array.
 static SmallVector<int64_t> getPositionFromIndices(ArrayRef<unsigned> indices) {
   SmallVector<int64_t> position;
   llvm::append_range(position, indices);
   return position;
 }
 
+/// Drops the first `n` elements of the `values` array.
+static SmallVector<llvm::Value *> dropFront(ArrayRef<llvm::Value *> values,
+                                            int64_t n) {
+  SmallVector<llvm::Value *> result(values.drop_front(n));
+  return result;
+}
+
 DataLayoutSpecInterface
 mlir::translateDataLayout(const llvm::DataLayout &dataLayout,
                           MLIRContext *context) {
   assert(context && "expected MLIR context");
   std::string layoutstr = dataLayout.getStringRepresentation();
 
   // Remaining unhandled default layout defaults
   // e (little endian if not set)

mlir/test/Target/LLVMIR/Import/intrinsic.ll

 
 ; CHECK-LABEL: llvm.func @get_active_lane_mask
 define <7 x i1> @get_active_lane_mask(i64 %0, i64 %1) {
   ; CHECK: llvm.intr.get.active.lane.mask %{{.*}}, %{{.*}} : i64, i64 to vector<7xi1>
   %3 = call <7 x i1> @llvm.get.active.lane.mask.v7i1.i64(i64 %0, i64 %1)
   ret <7 x i1> %3
 }
 
-; TODO: masked load store intrinsics should be handled specially.
-define void @masked_load_store_intrinsics(<7 x float>* %0, <7 x i1> %1) {
-  %3 = call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %0, i32 1, <7 x i1> %1, <7 x float> undef)
-  %4 = call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %0, i32 1, <7 x i1> %1, <7 x float> %3)
-  call void @llvm.masked.store.v7f32.p0v7f32(<7 x float> %4, <7 x float>* %0, i32 1, <7 x i1> %1)
+; CHECK-LABEL: @masked_load_store_intrinsics
+; CHECK-SAME:  %[[VEC:[a-zA-Z0-9]+]]
+; CHECK-SAME:  %[[MASK:[a-zA-Z0-9]+]]
+define void @masked_load_store_intrinsics(<7 x float>* %vec, <7 x i1> %mask) {
+  ; CHECK:  %[[UNDEF:.+]] = llvm.mlir.undef
+  ; CHECK:  %[[VAL1:.+]] = llvm.intr.masked.load %[[VEC]], %[[MASK]], %[[UNDEF]] {alignment = 1 : i32}
+  ; CHECK-SAME:  (!llvm.ptr<vector<7xf32>>, vector<7xi1>, vector<7xf32>) -> vector<7xf32>
+  %1 = call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %vec, i32 1, <7 x i1> %mask, <7 x float> undef)
+  ; CHECK:  %[[VAL2:.+]] = llvm.intr.masked.load %[[VEC]], %[[MASK]], %[[VAL1]] {alignment = 4 : i32}
+  %2 = call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %vec, i32 4, <7 x i1> %mask, <7 x float> %1)
+  ; CHECK:  llvm.intr.masked.store %[[VAL2]], %[[VEC]], %[[MASK]] {alignment = 8 : i32}
+  ; CHECK-SAME:  vector<7xf32>, vector<7xi1> into !llvm.ptr<vector<7xf32>>
+  call void @llvm.masked.store.v7f32.p0v7f32(<7 x float> %2, <7 x float>* %vec, i32 8, <7 x i1> %mask)
   ret void
 }
 
-; TODO: masked gather scatter intrinsics should be handled specially.
-define void @masked_gather_scatter_intrinsics(<7 x float*> %0, <7 x i1> %1) {
-  %3 = call <7 x float> @llvm.masked.gather.v7f32.v7p0f32(<7 x float*> %0, i32 1, <7 x i1> %1, <7 x float> undef)
-  %4 = call <7 x float> @llvm.masked.gather.v7f32.v7p0f32(<7 x float*> %0, i32 1, <7 x i1> %1, <7 x float> %3)
-  call void @llvm.masked.scatter.v7f32.v7p0f32(<7 x float> %4, <7 x float*> %0, i32 1, <7 x i1> %1)
+; CHECK-LABEL: @masked_gather_scatter_intrinsics
+; CHECK-SAME:  %[[VEC:[a-zA-Z0-9]+]]
+; CHECK-SAME:  %[[MASK:[a-zA-Z0-9]+]]
+define void @masked_gather_scatter_intrinsics(<7 x float*> %vec, <7 x i1> %mask) {
+  ; CHECK:  %[[UNDEF:.+]] = llvm.mlir.undef
+  ; CHECK:  %[[VAL1:.+]] = llvm.intr.masked.gather %[[VEC]], %[[MASK]], %[[UNDEF]] {alignment = 1 : i32}
+  ; CHECK-SAME:  (!llvm.vec<7 x ptr<f32>>, vector<7xi1>, vector<7xf32>) -> vector<7xf32>
+  %1 = call <7 x float> @llvm.masked.gather.v7f32.v7p0f32(<7 x float*> %vec, i32 1, <7 x i1> %mask, <7 x float> undef)
+  ; CHECK:  %[[VAL2:.+]] = llvm.intr.masked.gather %[[VEC]], %[[MASK]], %[[VAL1]] {alignment = 4 : i32}
+  %2 = call <7 x float> @llvm.masked.gather.v7f32.v7p0f32(<7 x float*> %vec, i32 4, <7 x i1> %mask, <7 x float> %1)
+  ; CHECK:  llvm.intr.masked.scatter %[[VAL2]], %[[VEC]], %[[MASK]] {alignment = 8 : i32}
+  ; CHECK-SAME:  vector<7xf32>, vector<7xi1> into !llvm.vec<7 x ptr<f32>>
+  call void @llvm.masked.scatter.v7f32.v7p0f32(<7 x float> %2, <7 x float*> %vec, i32 8, <7 x i1> %mask)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @masked_expand_compress_intrinsics
 define void @masked_expand_compress_intrinsics(float* %0, <7 x i1> %1, <7 x float> %2) {
   ; CHECK: %[[val1:.+]] = "llvm.intr.masked.expandload"(%{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<f32>, vector<7xi1>, vector<7xf32>) -> vector<7xf32>
   %4 = call <7 x float> @llvm.masked.expandload.v7f32(float* %0, <7 x i1> %1, <7 x float> %2)
   ; CHECK: "llvm.intr.masked.compressstore"(%[[val1]], %{{.*}}, %{{.*}}) : (vector<7xf32>, !llvm.ptr<f32>, vector<7xi1>) -> ()

mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp

static StringLoc findNextVariable(StringRef str) {
   size_t endPos = str.find_if_not([](char c) { return isAlnum(c) || c == '_'; },
                                   startPos + 1);
   if (endPos == std::string::npos)
     endPos = str.size();
 
   return {startPos, endPos - startPos};
 }
 
-// Check if `name` is the name of the variadic operand of `op`.  The variadic
-// operand can only appear at the last position in the list of operands.
+// Check if `name` is a variadic operand of `op`. Seach all operands since the
+// MLIR and LLVM IR operand order may differ and only for the latter the
+// variadic operand is guaranteed to be at the end of the operands list.
 static bool isVariadicOperandName(const tblgen::Operator &op, StringRef name) {
-  unsigned numOperands = op.getNumOperands();
-  if (numOperands == 0)
-    return false;
-  const auto &operand = op.getOperand(numOperands - 1);
-  return operand.isVariableLength() && operand.name == name;
+  for (int i = 0, e = op.getNumOperands(); i < e; ++i)
+    if (op.getOperand(i).name == name)
+      return op.getOperand(i).isVariableLength();
 }
 
 // Check if `result` is a known name of a result of `op`.
 static bool isResultName(const tblgen::Operator &op, StringRef name) {
   for (int i = 0, e = op.getNumResults(); i < e; ++i)
     if (op.getResultName(i) == name)
       return true;
   return false;
while (StringLoc loc = findNextVariable(builderStrRef)) {
     auto name = loc.in(builderStrRef).drop_front();
     // First, insert the non-matched part as is.
     bs << builderStrRef.substr(0, loc.pos);
     // Then, rewrite the name based on its kind.
     FailureOr<int> argIndex = getArgumentIndex(op, name);
     if (succeeded(argIndex)) {
       // Access the LLVM IR operand that maps to the given argument index using
       // the provided argument indices mapping.
-      // FIXME: support trailing variadic arguments.
-      int64_t operandIdx = llvmArgIndices[*argIndex];
-      if (operandIdx < 0)
+      int64_t idx = llvmArgIndices[*argIndex];
+      if (idx < 0)
         PrintFatalError(record.getLoc(), "expected non-negative operand index");
-      assert(!isVariadicOperandName(op, name) && "unexpected variadic operand");
-      bs << formatv("processValue(llvmOperands[{0}])", operandIdx);
+      bool isVariadicOperand = isVariadicOperandName(op, name);
+      auto result =
+          isVariadicOperand
+              ? formatv("processValues(dropFront(llvmOperands, {0}))", idx)
+              : formatv("processValue(llvmOperands[{0}])", idx);
+      bs << result;
     } else if (isResultName(op, name)) {
       if (op.getNumResults() != 1)
         PrintFatalError(record.getLoc(), "expected op to have one result");
       bs << formatv("mapValue(inst)");
     } else if (name == "_int_attr") {
       bs << "matchIntegerAttr";
     } else if (name == "_resultType") {
       bs << "convertType(inst->getType())";