[mlir] Support opaque types in LLVM IR -> MLIR translation

[ftynse]

Summary

LLVM IR is moving towards adoption of opaque pointer types. These require extra
information to be passed when constructing some operations, in particular GEP
and Alloca. Adapt the builders of said operations and modify the translation
code to handle both opaque and non-opaque pointers.

This incidentally adds the translation for Alloca alignment and fixes the translation
of struct-related GEP indices that must be constant.

[ftynse]

Rebase.

[wsmoses]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

876

See comment in test, but just because an index is a constant int doesn't mean that it's a struct or array offset (though those require a constant integer)

mlir/test/Target/LLVMIR/import.ll

26

Does this actually make sense to have as the constant attribute, since its not a struct or array index?

[ftynse]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

876

The relationship is not a bijection. Struct indices must be static constants, other indices may or may not be.

mlir/test/Target/LLVMIR/import.ll

26

It shouldn't be an issue.

[wsmoses]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

876

I suppose that's fair. That said, I do feel like it is beneficial to not promote non struct/array constants, if possible (and even if not semantically required).

Specifically, I recall the reason we introduced the attribute is to prevent inadvertant CSE where two geps with constant bounds from different branch predecessors were turned into one GEP with a new block argument. This would be illegal since LLVM needs the fixed offset. For non struct/array constants this wasn't needed as an attribute.

If there's no other reasons for the GEP to have a constant that cropped up since, it would be nice to preserve the offset as an op specifically to enable those optimizations, since we know it's legal.

I'm okay (though not over-the-moon) with the change in the PR, but since it is changing the semantics of importing GEP, it may be worthwhile considering the implications.

[ftynse]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

876

Fair enough.

[ftynse]

Address review.

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

llvm::Type *elementType = moduleTranslation.convertType(
         $elem_type ? *$elem_type
                    : op.getType().cast<LLVMPointerType>().getElementType());
     auto *inst = builder.CreateAlloca(elementType, addrSpace, $arraySize);
     }] # setAlignmentCode # [{
     $res = inst;
   }];
   let builders = [
     OpBuilder<(ins "Type":$resultType, "Value":$arraySize,
                "unsigned":$alignment),
     [{
+      assert(!resultType.cast<LLVMPointerType>().isOpaque() &&
+             "pass the allocated type explicitly if opaque pointers are used");
       if (alignment == 0)
         return build($_builder, $_state, resultType, arraySize, IntegerAttr(),
                      TypeAttr());
       build($_builder, $_state, resultType, arraySize,
         $_builder.getI64IntegerAttr(alignment), TypeAttr());
-  }]>];
+    }]>,
+    OpBuilder<(ins "Type":$resultType, "Type":$elementType, "Value":$arraySize,
+               CArg<"unsigned", "0">:$alignment),
+    [{
+      TypeAttr elemTypeAttr =
+          resultType.cast<LLVMPointerType>().isOpaque() ?
+          TypeAttr::get(elementType) : TypeAttr();
+      build($_builder, $_state, resultType, arraySize,
+            alignment == 0 ? IntegerAttr()
+                           : $_builder.getI64IntegerAttr(alignment),
+            elemTypeAttr);
+
+    }]>
+    ];
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_GEPOp : LLVM_Op<"getelementptr", [NoSideEffect]> {
   let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$base,
                    Variadic<LLVM_ScalarOrVectorOf<AnyInteger>>:$indices,
                    I32ElementsAttr:$structIndices,
                    OptionalAttr<TypeAttr>:$elem_type);
   let results = (outs LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$res);
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins "Type":$resultType, "Value":$basePtr, "ValueRange":$indices,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
     OpBuilder<(ins "Type":$resultType, "Value":$basePtr, "ValueRange":$indices,
                "ArrayRef<int32_t>":$structIndices,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "Type":$resultType, "Type":$basePtrType, "Value":$basePtr,
+               "ValueRange":$indices, "ArrayRef<int32_t>":$structIndices,
+               CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
   ];
   let llvmBuilder = [{
     SmallVector<llvm::Value *> indices;
     indices.reserve($structIndices.size());
     unsigned operandIdx = 0;
     for (int32_t structIndex : $structIndices.getValues<int32_t>()) {
       if (structIndex == GEPOp::kDynamicIndex)
         indices.push_back($indices[operandIdx++]);
       else
         indices.push_back(builder.getInt32(structIndex));
     }
     Type baseElementType = op.getSourceElementType();
     llvm::Type *elementType = moduleTranslation.convertType(baseElementType);
     $res = builder.CreateGEP(elementType, $base, indices);
   }];
   let assemblyFormat = [{
     $base `[` custom<GEPIndices>($indices, $structIndices) `]` attr-dict
     `:` functional-type(operands, results) (`,` $elem_type^)?
   }];
 
   let extraClassDeclaration = [{
     constexpr static int kDynamicIndex = std::numeric_limits<int32_t>::min();
+
+    /// Returns the type pointed to by the pointer argument of this GEP.
     Type getSourceElementType();
+
+    /// Populates `indices` with positions of GEP indices that correspond to
+    /// LLVMStructTypes potentially nested in the given `sourceElementType`,
+    /// which is the type pointed to by the pointer argument of a GEP. If
+    /// `structSizes` is provided, it is populated with sizes of the indexed
+    /// structs for bounds verification purposes.
+    static void findKnownStructIndices(
+        Type sourceElementType, SmallVectorImpl<unsigned> &indices,
+        SmallVectorImpl<unsigned> *structSizes = nullptr);
+
   }];
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_LoadOp : LLVM_Op<"load">, MemoryOpWithAlignmentAndAttributes {
   let arguments = (ins Arg<LLVM_PointerTo<LLVM_LoadableType>, "", [MemRead]>:$addr,
                    OptionalAttr<SymbolRefArrayAttr>:$access_groups,

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

void AllocaOp::print(OpAsmPrinter &p) {
   if (!elemTy)
     elemTy = *getElemType();
 
   auto funcTy =
       FunctionType::get(getContext(), {getArraySize().getType()}, {getType()});
 
   p << ' ' << getArraySize() << " x " << elemTy;
   if (getAlignment().hasValue() && *getAlignment() != 0)
-    p.printOptionalAttrDict((*this)->getAttrs());
+    p.printOptionalAttrDict((*this)->getAttrs(), {kElemTypeAttrName});
   else
     p.printOptionalAttrDict((*this)->getAttrs(),
                             {"alignment", kElemTypeAttrName});
   p << " : " << funcTy;
 }
 
 // <operation> ::= `llvm.alloca` ssa-use `x` type attribute-dict?
 //                 `:` type `,` type
llvm::TypeSwitch<Type>(type)
       });
 }
 
 /// Populates `indices` with positions of GEP indices that correspond to
 /// LLVMStructTypes potentially nested in the given `baseGEPType`, which must
 /// be either an LLVMPointer type or a vector thereof. If `structSizes` is
 /// provided, it is populated with sizes of the indexed structs for bounds
 /// verification purposes.
-static void
-findKnownStructIndices(Type sourceElementType,
+void GEPOp::findKnownStructIndices(Type sourceElementType,
                                    SmallVectorImpl<unsigned> &indices,
-                       SmallVectorImpl<unsigned> *structSizes = nullptr) {
+                                   SmallVectorImpl<unsigned> *structSizes) {
   SmallPtrSet<Type, 4> visited;
   recordStructIndices(sourceElementType, /*currentIndex=*/1, indices,
                       structSizes, visited);
 }
 
 void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
                   Value basePtr, ValueRange operands,
                   ArrayRef<NamedAttribute> attributes) {
if (auto fixedVectorType = type.dyn_cast<LLVMFixedVectorType>())
     return fixedVectorType.getElementType();
   return type;
 }
 
 void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
                   Value basePtr, ValueRange indices,
                   ArrayRef<int32_t> structIndices,
                   ArrayRef<NamedAttribute> attributes) {
+  auto ptrType =
+      extractVectorElementType(basePtr.getType()).cast<LLVMPointerType>();
+  assert(!ptrType.isOpaque() &&
+         "expected non-opaque pointer, provide elementType explicitly when "
+         "opaque pointers are used");
+  build(builder, result, resultType, ptrType.getElementType(), basePtr, indices,
+        structIndices, attributes);
+}
+
+void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
+                  Type elementType, Value basePtr, ValueRange indices,
+                  ArrayRef<int32_t> structIndices,
+                  ArrayRef<NamedAttribute> attributes) {
   SmallVector<Value> remainingIndices;
   SmallVector<int32_t> updatedStructIndices(structIndices.begin(),
                                             structIndices.end());
   SmallVector<unsigned> structRelatedPositions;
-  auto ptrType =
-      extractVectorElementType(basePtr.getType()).cast<LLVMPointerType>();
-  assert(!ptrType.isOpaque() && "expected non-opaque pointer");
-  findKnownStructIndices(ptrType.getElementType(), structRelatedPositions);
+  findKnownStructIndices(elementType, structRelatedPositions);
 
   SmallVector<unsigned> operandsToErase;
   for (unsigned pos : structRelatedPositions) {
     // GEP may not be indexing as deep as some structs are located.
     if (pos >= structIndices.size())
       continue;
 
     // If the index is already static, it's fine.
void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
   assert(remainingIndices.size() == static_cast<size_t>(llvm::count(
                                         updatedStructIndices, kDynamicIndex)) &&
          "expected as many index operands as dynamic index attr elements");
 
   result.addTypes(resultType);
   result.addAttributes(attributes);
   result.addAttribute("structIndices",
                       builder.getI32TensorAttr(updatedStructIndices));
+  if (extractVectorElementType(basePtr.getType())
+          .cast<LLVMPointerType>()
+          .isOpaque())
+    result.addAttribute(kElemTypeAttrName, TypeAttr::get(elementType));
   result.addOperands(basePtr);
   result.addOperands(remainingIndices);
 }
 
 static ParseResult
 parseGEPIndices(OpAsmParser &parser,
                 SmallVectorImpl<OpAsmParser::UnresolvedOperand> &indices,
                 DenseIntElementsAttr &structIndices) {

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

LogicalResult Importer::processInstruction(llvm::Instruction *inst) {
   case llvm::Instruction::SRem:
   case llvm::Instruction::FRem:
   case llvm::Instruction::Shl:
   case llvm::Instruction::LShr:
   case llvm::Instruction::AShr:
   case llvm::Instruction::And:
   case llvm::Instruction::Or:
   case llvm::Instruction::Xor:
-  case llvm::Instruction::Alloca:
   case llvm::Instruction::Load:
   case llvm::Instruction::Store:
   case llvm::Instruction::Ret:
   case llvm::Instruction::Resume:
   case llvm::Instruction::Trunc:
   case llvm::Instruction::ZExt:
   case llvm::Instruction::SExt:
   case llvm::Instruction::FPToUI:
if (!inst->getType()->isVoidTy()) {
         return failure();
       state.addTypes(type);
     }
     Operation *op = b.create(state);
     if (!inst->getType()->isVoidTy())
       v = op->getResult(0);
     return success();
   }
+  case llvm::Instruction::Alloca: {
+    Value size = processValue(inst->getOperand(0));
+    if (!size)
+      return failure();
+
+    auto *allocaInst = cast<llvm::AllocaInst>(inst);
+    v = b.create<AllocaOp>(loc, processType(inst->getType()),
+                           processType(allocaInst->getAllocatedType()), size,
+                           allocaInst->getAlign().value());
+    return success();
+  }
   case llvm::Instruction::ICmp: {
     Value lhs = processValue(inst->getOperand(0));
     Value rhs = processValue(inst->getOperand(1));
     if (!lhs || !rhs)
       return failure();
     v = b.create<ICmpOp>(
         loc, getICmpPredicate(cast<llvm::ICmpInst>(inst)->getPredicate()), lhs,
         rhs);
case llvm::Instruction::Fence: {
     }
     b.create<FenceOp>(loc, getLLVMAtomicOrdering(fence->getOrdering()),
                       syncscope);
     return success();
   }
   case llvm::Instruction::GetElementPtr: {
     // FIXME: Support inbounds GEPs.
     llvm::GetElementPtrInst *gep = cast<llvm::GetElementPtrInst>(inst);
-    SmallVector<Value, 4> ops;
-    for (auto *op : gep->operand_values()) {
-      Value value = processValue(op);
-      if (!value)
+    Value basePtr = processValue(gep->getOperand(0));
+    SmallVector<int32_t> staticIndices;
+    SmallVector<Value> dynamicIndices;
+    Type sourceElementType = processType(gep->getSourceElementType());
+    SmallVector<unsigned> staticIndexPositions;
+    GEPOp::findKnownStructIndices(sourceElementType, staticIndexPositions);
+
+    for (const auto &en :
+         llvm::enumerate(llvm::drop_begin(gep->operand_values()))) {
+      llvm::Value *operand = en.value();
+      if (llvm::find(staticIndexPositions, en.index()) ==
+          staticIndexPositions.end()) {
+        staticIndices.push_back(GEPOp::kDynamicIndex);
+        dynamicIndices.push_back(processValue(operand));
+        if (!dynamicIndices.back())
           return failure();
-      ops.push_back(value);
+      } else {
+        auto *constantInt = cast<llvm::ConstantInt>(operand);
+        staticIndices.push_back(
+            static_cast<int32_t>(constantInt->getValue().getZExtValue()));
       }
+    }
+
     Type type = processType(inst->getType());
     if (!type)
       return failure();
-    v = b.create<GEPOp>(loc, type, ops[0],
-                        llvm::makeArrayRef(ops).drop_front());
+    v = b.create<GEPOp>(loc, type, sourceElementType, basePtr, dynamicIndices,
+                        staticIndices);
     return success();
   }
   }
 }
 
 FlatSymbolRefAttr Importer::getPersonalityAsAttr(llvm::Function *f) {
   if (!f->hasPersonalityFn())
     return nullptr;

mlir/lib/Target/LLVMIR/TypeFromLLVM.cpp

private:
 
   /// Translates the given integer type.
   Type translate(llvm::IntegerType *type) {
     return IntegerType::get(&context, type->getBitWidth());
   }
 
   /// Translates the given pointer type.
   Type translate(llvm::PointerType *type) {
+    if (type->isOpaque())
+      return LLVM::LLVMPointerType::get(&context, type->getAddressSpace());
+
     return LLVM::LLVMPointerType::get(
-        translateType(type->getPointerElementType()), type->getAddressSpace());
+        translateType(type->getNonOpaquePointerElementType()),
+        type->getAddressSpace());
   }
 
   /// Translates the given structure type.
   Type translate(llvm::StructType *type) {
     SmallVector<Type, 8> subtypes;
     if (type->isLiteral()) {
       translateTypes(type->subtypes(), subtypes);
       return LLVM::LLVMStructType::getLiteral(&context, subtypes,

mlir/test/Target/LLVMIR/import-opaque.ll

+; RUN: mlir-translate -import-llvm -split-input-file -opaque-pointers %s | FileCheck %s
+
+; CHECK-LABEL: @opaque_ptr_load
+define i32 @opaque_ptr_load(ptr %0) {
+  ; CHECK: = llvm.load %{{.*}} : !llvm.ptr -> i32
+  %2 = load i32, ptr %0, align 4
+  ret i32 %2
+}
+
+; // -----
+
+; CHECK-LABEL: @opaque_ptr_store
+define void @opaque_ptr_store(i32 %0, ptr %1) {
+  ; CHECK: llvm.store %{{.*}}, %{{.*}} : i32, !llvm.ptr
+  store i32 %0, ptr %1, align 4
+  ret void
+}
+
+; // -----
+
+; CHECK-LABEL: @opaque_ptr_ptr_store
+define void @opaque_ptr_ptr_store(ptr %0, ptr %1) {
+  ; CHECK: llvm.store %{{.*}}, %{{.*}} : !llvm.ptr, !llvm.ptr
+  store ptr %0, ptr %1, align 8
+  ret void
+}
+
+; // -----
+
+; CHECK-LABEL: @opaque_ptr_alloca
+define ptr @opaque_ptr_alloca(i32 %0) {
+  ; CHECK: = llvm.alloca %{{.*}} x f32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
+  %2 = alloca float, i32 %0, align 4
+  ret ptr %2
+}
+
+; // -----
+
+; CHECK-LABEL: @opaque_ptr_gep
+define ptr @opaque_ptr_gep(ptr %0, i32 %1) {
+  ; CHECK: = llvm.getelementptr %{{.*}}[%{{.*}}] : (!llvm.ptr, i32) -> !llvm.ptr, f32
+  %3 = getelementptr float, ptr %0, i32 %1
+  ret ptr %3
+}
+
+; // -----
+
+; CHECK-LABEL: @opaque_ptr_gep
+define ptr @opaque_ptr_gep_struct(ptr %0, i32 %1){
+  ; CHECK: = llvm.getelementptr %{{.*}}[%{{.*}}, 0, 1] : (!llvm.ptr, i32) -> !llvm.ptr, !llvm.struct<(struct<(f32, f64)>, struct<(i32, i64)>)>
+  %3 = getelementptr { { float, double }, { i32, i64 } }, ptr %0, i32 %1, i32 0, i32 1
+  ret ptr %3
+}

mlir/test/Target/LLVMIR/import.ll

 
 ; CHECK: llvm.mlir.global private @alig64(42 : i64) {alignment = 64 : i64, dso_local} : i64
 @alig64 = private global i64 42, align 64
 
 @g4 = external global i32, align 8
 ; CHECK: llvm.mlir.global internal constant @int_gep() {dso_local} : !llvm.ptr<i32> {
 ; CHECK-DAG:   %[[addr:[0-9]+]] = llvm.mlir.addressof @g4 : !llvm.ptr<i32>
 ; CHECK-DAG:   %[[c2:[0-9]+]] = llvm.mlir.constant(2 : i32) : i32
 ; CHECK-NEXT:  %[[gepinit:[0-9]+]] = llvm.getelementptr %[[addr]][%[[c2]]] : (!llvm.ptr<i32>, i32) -> !llvm.ptr<i32>
 ; CHECK-NEXT:  llvm.return %[[gepinit]] : !llvm.ptr<i32>
 ; CHECK-NEXT: }
 @int_gep = internal constant i32* getelementptr (i32, i32* @g4, i32 2)
 
 ;
 ; dso_local attribute
 ;
 

 @_ZTIii= external dso_local constant i8**
 declare void @foo(i8*)
 declare i8* @bar(i8*)
 declare i32 @__gxx_personality_v0(...)
 
 ; CHECK-LABEL: @invokeLandingpad
 define i32 @invokeLandingpad() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
   ; CHECK: %[[a1:[0-9]+]] = llvm.bitcast %{{[0-9]+}} : !llvm.ptr<ptr<ptr<i8>>> to !llvm.ptr<i8>
-  ; CHECK: %[[a3:[0-9]+]] = llvm.alloca %{{[0-9]+}} x i8 : (i32) -> !llvm.ptr<i8>
+  ; CHECK: %[[a3:[0-9]+]] = llvm.alloca %{{[0-9]+}} x i8 {alignment = 1 : i64} : (i32) -> !llvm.ptr<i8>
   %1 = alloca i8
   ; CHECK: llvm.invoke @foo(%[[a3]]) to ^bb2 unwind ^bb1 : (!llvm.ptr<i8>) -> ()
   invoke void @foo(i8* %1) to label %4 unwind label %2
 
 ; CHECK: ^bb1:
   ; CHECK: %{{[0-9]+}} = llvm.landingpad (catch %{{[0-9]+}} : !llvm.ptr<ptr<i8>>) (catch %[[a1]] : !llvm.ptr<i8>) (filter %{{[0-9]+}} : !llvm.array<1 x i8>) : !llvm.struct<(ptr<i8>, i32)>
   %3 = landingpad { i8*, i32 } catch i8** @_ZTIi catch i8* bitcast (i8*** @_ZTIii to i8*)
   ; FIXME: Change filter to a constant array once they are handled.