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

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.

Reviewed By: bkramer

Differential Revision: https://reviews.llvm.org/D118006

Author: nicolasvasilache

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

@@ -1934,17 +1934,25 @@ def LLVM_InlineAsmOp : LLVM_Op<"inline_asm", []> {
         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

@@ -37,7 +37,10 @@ Value mlir::x86vector::avx2::inline_asm::mm256BlendPsAsm(
   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);
 }
 

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

@@ -330,11 +330,32 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
                                    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();
   }
 

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

@@ -0,0 +1,49 @@
+// 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
+  }
+}
+