[mlir][LLVM] Add support for operand_attrs to InlineAsmOp

[nicolasvasilache]

Summary

This revision adds enough support to allow InlineAsmOp to work properly with indirect memory constraints "*m".
These require an explicit "elementtype" TypeAttr on the operands to pass LLVM verification and need to be provided.

[bkramer]

Looks good

[nicolasvasilache]

Rebase and fix condition.

[nicolasvasilache]

Investigate post revert

[nicolasvasilache]

build is green resubmitting.

[nicolasvasilache]

@mehdi_amini what triggered the previous revert?

[nicolasvasilache]

ok I see it now in https://buildkite.com/mlir/mlir-core/builds/20054#49efb465-9bb6-483e-8e92-4ab92560788d, looking

[nicolasvasilache]

This seems related to LLI + Intel SDE running through mlir-check.
Sent a patch that temporarily disables the LLI part of the test, the mlir-opt + mlir-translate is still a meaningful test: c2fa5ff5adba61dc5b0d477c0af1ea00ec2c596a.

Help to understand what is happening here would be appreciated.

[nicolasvasilache]

This is now fixed properly in 31b7f8857f4cd9c182b397ebf55db049318b3d7d: the test was not returning and would silently return non-0.

Thanks @bkramer for the correct hypothesis!

[mehdi_amini]

In D118006#3286904, @nicolasvasilache wrote:

@mehdi_amini what triggered the previous revert?

Sorry, I forgot to come back here to past the link! Glad you figured it out :)

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

def LLVM_InlineAsmOp : LLVM_Op<"inline_asm", []> {
   }];
   let arguments = (
     ins Variadic<LLVM_Type>:$operands,
         StrAttr:$asm_string,
         StrAttr:$constraints,
         UnitAttr:$has_side_effects,
         UnitAttr:$is_align_stack,
         OptionalAttr<
-          DefaultValuedAttr<AsmATTOrIntel, "AsmDialect::AD_ATT">>:$asm_dialect);
+          DefaultValuedAttr<AsmATTOrIntel, "AsmDialect::AD_ATT">>:$asm_dialect,
+        OptionalAttr<ArrayAttr>:$operand_attrs);
 
   let results = (outs Optional<LLVM_Type>:$res);
 
   let assemblyFormat = [{
     (`has_side_effects` $has_side_effects^)?
     (`is_align_stack` $is_align_stack^)?
     (`asm_dialect` `=` $asm_dialect^)?
+    (`operand_attrs` `=` $operand_attrs^)?
     attr-dict
     $asm_string `,` $constraints
     operands `:` functional-type(operands, results)
    }];
+
+  let extraClassDeclaration = [{
+    static StringRef getElementTypeAttrName() {
+      return "elementtype";
+    }
+  }];
 }
 #endif // LLVMIR_OPS

mlir/lib/Dialect/X86Vector/Transforms/AVXTranspose.cpp

Value mlir::x86vector::avx2::inline_asm::mm256BlendPsAsm(
   auto asmDialectAttr =
       LLVM::AsmDialectAttr::get(b.getContext(), LLVM::AsmDialect::AD_Intel);
   const auto *asmTp = "vblendps $0, $1, $2, {0}";
   const auto *asmCstr =
       "=x,x,x"; // Careful: constraint parser is very brittle: no ws!
   SmallVector<Value> asmVals{v1, v2};
   auto asmStr = llvm::formatv(asmTp, llvm::format_hex(mask, /*width=*/2)).str();
   auto asmOp = b.create<LLVM::InlineAsmOp>(
-      v1.getType(), asmVals, asmStr, asmCstr, false, false, asmDialectAttr);
+      v1.getType(), /*operands=*/asmVals, /*asm_string=*/asmStr,
+      /*constraints=*/asmCstr, /*has_side_effects=*/false,
+      /*is_align_stack=*/false, /*asm_dialect=*/asmDialectAttr,
+      /*operand_attrs=*/ArrayAttr());
   return asmOp.getResult(0);
 }
 
 Value mlir::x86vector::avx2::intrin::mm256UnpackLoPs(ImplicitLocOpBuilder &b,
                                                      Value v1, Value v2) {
   return b.create<vector::ShuffleOp>(
       v1, v2, ArrayRef<int64_t>{0, 8, 1, 9, 4, 12, 5, 13});
 }

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

llvm::InlineAsm *inlineAsmInst =
                   inlineAsmOp.getIsAlignStack(),
                   convertAsmDialectToLLVM(*inlineAsmOp.getAsmDialect()))
             : llvm::InlineAsm::get(static_cast<llvm::FunctionType *>(
                                        moduleTranslation.convertType(ft)),
                                    inlineAsmOp.getAsmString(),
                                    inlineAsmOp.getConstraints(),
                                    inlineAsmOp.getHasSideEffects(),
                                    inlineAsmOp.getIsAlignStack());
-    llvm::Value *result = builder.CreateCall(
+    llvm::CallInst *inst = builder.CreateCall(
         inlineAsmInst,
         moduleTranslation.lookupValues(inlineAsmOp.getOperands()));
+    if (auto maybeOperandAttrs = inlineAsmOp.getOperandAttrs()) {
+      llvm::AttributeList attrList;
+      for (const auto &it : llvm::enumerate(*maybeOperandAttrs)) {
+        Attribute attr = it.value();
+        if (!attr)
+          continue;
+        DictionaryAttr dAttr = attr.cast<DictionaryAttr>();
+        TypeAttr tAttr =
+            dAttr.get(InlineAsmOp::getElementTypeAttrName()).cast<TypeAttr>();
+        llvm::AttrBuilder b(moduleTranslation.getLLVMContext());
+        llvm::Type *ty = moduleTranslation.convertType(tAttr.getValue());
+        b.addTypeAttr(llvm::Attribute::ElementType, ty);
+        // shift to account for the returned value (this is always 1 aggregate
+        // value in LLVM).
+        int shift = (opInst.getNumResults() > 0) ? 1 : 0;
+        attrList = attrList.addAttributesAtIndex(
+            moduleTranslation.getLLVMContext(), it.index() + shift, b);
+      }
+      inst->setAttributes(attrList);
+    }
+
     if (opInst.getNumResults() != 0)
-      moduleTranslation.mapValue(opInst.getResult(0), result);
+      moduleTranslation.mapValue(opInst.getResult(0), inst);
     return success();
   }
 
   if (auto invOp = dyn_cast<LLVM::InvokeOp>(opInst)) {
     auto operands = moduleTranslation.lookupValues(invOp.getCalleeOperands());
     ArrayRef<llvm::Value *> operandsRef(operands);
     llvm::Instruction *result;
     if (auto attr = opInst.getAttrOfType<FlatSymbolRefAttr>("callee")) {

mlir/test/Integration/Dialect/Vector/CPU/X86Vector/test-inline-asm-vector-avx512.mlir

+// RUN: mlir-opt %s -convert-linalg-to-loops -convert-vector-to-scf='full-unroll=true' -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-memref-to-llvm  -convert-std-to-llvm='use-bare-ptr-memref-call-conv=1' -convert-arith-to-llvm -reconcile-unrealized-casts |\
+// RUN: mlir-translate --mlir-to-llvmir |\
+// RUN: %lli --entry-function=entry --mattr="avx512f" | \
+// RUN: FileCheck %s
+
+module {
+
+  // printf format string "%i\n", char by char:   %    i  \n  0
+  llvm.mlir.global private @pct_i_newline(dense<[37, 105, 10, 0]> : tensor<4xi8>)
+    : !llvm.array<4xi8>
+  // an array of 16 i32 of values [0..15]
+  llvm.mlir.global private @const16(
+    dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16 x i32>)
+      : !llvm.array<16 x i32>
+
+  // declare void @printf(i8*, ...)
+  llvm.func @printf(!llvm.ptr<i8>, ...)
+
+  llvm.func @entry() {
+    %c0 = llvm.mlir.constant(0 : index) : i64
+
+    %1 = llvm.mlir.addressof @const16 : !llvm.ptr<array<16 x i32>>
+    %ptr = llvm.getelementptr %1[%c0, %c0]
+      : (!llvm.ptr<array<16 x i32>>, i64, i64) -> !llvm.ptr<i32>
+    %ptr2 = llvm.bitcast %ptr :  !llvm.ptr<i32> to !llvm.ptr<vector<16xi32>>
+    // operand_attrs of *m operands need to be piped through to LLVM for
+    // verification to pass.
+    %v = llvm.inline_asm
+        asm_dialect = intel
+        operand_attrs = [{ elementtype = vector<16xi32> }]
+        "vmovdqu32 $0, $1", "=x,*m" %ptr2
+      : (!llvm.ptr<vector<16xi32>>) -> vector<16xi32>
+
+    %2 = llvm.mlir.addressof @pct_i_newline : !llvm.ptr<array<4xi8>>
+    %ptrfmt = llvm.getelementptr %2[%c0, %c0]
+      : (!llvm.ptr<array<4xi8>>, i64, i64) -> !llvm.ptr<i8>
+
+    // CHECK: 0
+    %v0 = vector.extract %v[0]: vector<16xi32>
+    llvm.call @printf(%ptrfmt, %v0) : (!llvm.ptr<i8>, i32) -> ()
+
+    // CHECK: 9
+    %v9 = vector.extract %v[9]: vector<16xi32>
+    llvm.call @printf(%ptrfmt, %v9) : (!llvm.ptr<i8>, i32) -> ()
+
+    llvm.return
+  }
+}
+