Use opaque pointers when lowering to LLVMIR (#2190)

* Use opaque pointers when lowering to LLVMIR

Signed-off-by: Tung D. Le <[email protected]>

Author: tungld

docs/Testing.md

@@ -377,7 +377,7 @@ Again, these debug statements can then be activated by adding the `--debug-only=
 
 ## ONNX Model Zoo
 
-We provide a Python script [RunONNXModelZoo.py](../utils/RunONNXModelZoo.py) to check inference accuracy with models in the [ONNX model zoo](https://github.com/onnx/models).  [RunONNXModelZoo.py](../utils/RunONNXModelZoo.py) requires [RunONNXModel.py](../utils/RunONNXModel.py) to be in the same folder. For example,
+We provide a Python script [RunONNXModelZoo.py](../utils/RunONNXModelZoo.py) to check inference accuracy with models in the [ONNX model zoo](https://github.com/onnx/models).  [RunONNXModelZoo.py](../utils/RunONNXModelZoo.py) requires [RunONNXModel.py](../utils/RunONNXModel.py) to be in the same folder. For example, to check inference accuracy with mnist-8:
 
 ```bash
 $ mkdir test && cd test
@@ -388,3 +388,5 @@ $ ONNX_MLIR_HOME=/onnx-mlir/build/Release/ python RunONNXModelZoo.py -m mnist-8
 Run the script with `-h` to see all the options. In addition to the `-m` flag to specify a model and `-c` flag to specify the compile options, useful options are the `-k` flag to leave the onnx model in the current directory as a `.tgz` file, and the `-l debug` flag to print lots of debugging info.
 
 To find out which models are available, run the script with `-p` to print the list of available models; or `-m` followed by an incomplete name, and the script will suggest the exact names. 
+
+Without specifying a model using `-m`, the script will check all models in the ONNX model zoo.

src/Accelerators/NNPA/Conversion/ZLowToLLVM/CMakeLists.txt

@@ -7,6 +7,7 @@ add_onnx_mlir_library(OMZLowToLLVM
 
   LINK_LIBS PUBLIC
   MLIRLLVMCommonConversion
+  OMKrnlToLLVM
   OMLayoutHelper
   OMZLowOps
   OMMlirDialects

src/Accelerators/NNPA/Conversion/ZLowToLLVM/ZLowToLLVM.cpp

@@ -19,6 +19,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 
 #include "src/Accelerators/NNPA/Conversion/ZLowToLLVM/ZLowToLLVMCommon.hpp"
+#include "src/Conversion/KrnlToLLVM/KrnlToLLVMHelper.hpp"
 #include "src/Dialect/Mlir/DialectBuilder.hpp"
 #include "zdnn.h"
 
@@ -29,7 +30,7 @@ namespace zlow {
 
 ApiRegistry RegisterAllApis(MLIRContext *context) {
   auto voidTy = LLVM::LLVMVoidType::get(context);
-  auto opaquePtrTy = LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
+  auto opaquePtrTy = krnl::getI8PointerType(context);
   auto int32Ty = IntegerType::get(context, 32);
   auto int64Ty = IntegerType::get(context, 64);
 
@@ -105,7 +106,8 @@ Value ZTensorHelper::getPreTransformedDescPtr(zdnn_data_types zDNNDataType,
   Value one = create.llvm.constant(llvmI64Ty, (int64_t)1);
 
   Value preTransformedDescPtr = create.llvm._alloca(
-      LLVM::LLVMPointerType::get(llvmZTensorDescStructTy), one,
+      krnl::getPointerType(context, llvmZTensorDescStructTy),
+      llvmZTensorDescStructTy, one,
       /*alignment=*/0);
 
   // Prepare operands for calling the function that initializes the zTensor
@@ -145,7 +147,8 @@ Value ZTensorHelper::getTransformedDescPtr(
   Value one = create.llvm.constant(llvmI64Ty, (int64_t)1);
 
   Value transformedDescPtr = create.llvm._alloca(
-      LLVM::LLVMPointerType::get(llvmZTensorDescStructTy), one,
+      krnl::getPointerType(context, llvmZTensorDescStructTy),
+      llvmZTensorDescStructTy, one,
       /*alignment=*/0);
 
   if (isConcat) {
@@ -165,10 +168,11 @@ Value ZTensorHelper::getTransformedDescPtr(
 
 // Get the pointer to memref.
 Value ZTensorHelper::getAlignedI8Ptr(Value memRef) {
+  MLIRContext *context = rewriter.getContext();
   MultiDialectBuilder<LLVMBuilder> create(rewriter, loc);
   MemRefDescriptor descriptor(memRef);
   Value alignedPtr = descriptor.alignedPtr(rewriter, loc);
-  return create.llvm.bitcastI8Ptr(alignedPtr);
+  return create.llvm.bitcast(krnl::getI8PointerType(context), alignedPtr);
 }
 
 // Get buffer size from a transformed descriptor.
@@ -202,7 +206,8 @@ ZTensor ZTensorHelper::getZTensor(Value bufferPtr, zdnn_data_types dataType,
       getTransformedDescPtr(preTransformedDescPtr, isConcat, concatInfo);
   // Create the input zTensor.
   Value alloc =
-      create.llvm._alloca(LLVM::LLVMPointerType::get(llvmZTensorStructTy), one,
+      create.llvm._alloca(krnl::getPointerType(context, llvmZTensorStructTy),
+          llvmZTensorStructTy, one,
           /*alignment=*/0);
   // Buffer size.
   Value bufferSize = getBufferSize(transformedDescPtr);
@@ -235,7 +240,8 @@ ZTensor ZTensorHelper::getZTensor(Value preTransformedDescPtr,
   Type llvmZTensorStructTy = getZTensorStructTy(context);
   Value one = create.llvm.constant(rewriter.getI64Type(), (int64_t)1);
   Value alloc =
-      create.llvm._alloca(LLVM::LLVMPointerType::get(llvmZTensorStructTy), one,
+      create.llvm._alloca(krnl::getPointerType(context, llvmZTensorStructTy),
+          llvmZTensorStructTy, one,
           /*alignment=*/0);
   // clang-format off
   fillInZTensor(rewriter, loc, module, alloc,
@@ -370,10 +376,10 @@ std::vector<Value> getDimsFromShapeMemRefBySize(PatternRewriter &rewriter,
   Value alignedPtr = inputMRD.alignedPtr(rewriter, loc);
   Type int64Ty = IntegerType::get(context, 64);
   for (int64_t i = 0; i < size; ++i) {
-    Value index = create.llvm.constant(int64Ty, i);
-    Value alignedGep = create.llvm.getElemPtr(
-        LLVM::LLVMPointerType::get(int64Ty), alignedPtr, {index});
-    Value dimI64 = create.llvm.load(alignedGep);
+    Value alignedGep =
+        create.llvm.getElemPtr(krnl::getPointerType(context, int64Ty), int64Ty,
+            alignedPtr, ArrayRef<LLVM::GEPArg>{(int32_t)i});
+    Value dimI64 = create.llvm.load(int64Ty, alignedGep);
     dims.emplace_back(dimI64);
   }
   return dims;
@@ -462,16 +468,16 @@ Type getZTensorStructTy(MLIRContext *context) {
   Type llvmI1Ty = IntegerType::get(context, 1);
   Type llvmI8Ty = IntegerType::get(context, 8);
   Type llvmArrayI8Ty = LLVM::LLVMArrayType::get(llvmI8Ty, 32);
-  Type llvmI8PtrTy = LLVM::LLVMPointerType::get(llvmI8Ty);
+  Type llvmI8PtrTy = krnl::getPointerType(context, llvmI8Ty);
   Type llvmZTensorDescStructTy = getZTensorDescStructTy(context);
 
   SmallVector<Type, 4> zTensorTypeElements;
   // A pointer to pre-transformed descriptor struct type
   zTensorTypeElements.emplace_back(
-      LLVM::LLVMPointerType::get(llvmZTensorDescStructTy));
+      krnl::getPointerType(context, llvmZTensorDescStructTy));
   // A pointer to transformed descriptor struct type
   zTensorTypeElements.emplace_back(
-      LLVM::LLVMPointerType::get(llvmZTensorDescStructTy));
+      krnl::getPointerType(context, llvmZTensorDescStructTy));
   // zTensor size in bytes
   zTensorTypeElements.emplace_back(llvmI64Ty);
   // pointer to the zTensor in memory
@@ -490,8 +496,9 @@ Type getZTensorStructTy(MLIRContext *context) {
 /// Function to cast an LLVM pointer to an opaque LLVM pointer.
 Value toOpaquePtr(
     PatternRewriter &rewriter, Location loc, ModuleOp module, Value ptr) {
+  MLIRContext *context = rewriter.getContext();
   MultiDialectBuilder<LLVMBuilder> create(rewriter, loc);
-  return create.llvm.bitcastI8Ptr(ptr);
+  return create.llvm.bitcast(krnl::getI8PointerType(context), ptr);
 }
 
 void fillInZTensor(PatternRewriter &rewriter, Location loc, ModuleOp module,
@@ -501,48 +508,35 @@ void fillInZTensor(PatternRewriter &rewriter, Location loc, ModuleOp module,
   MultiDialectBuilder<LLVMBuilder> create(rewriter, loc);
 
   Type llvmI1Ty = IntegerType::get(context, 1);
-  Type llvmI8Ty = IntegerType::get(context, 8);
-  Type llvmI8PtrTy = LLVM::LLVMPointerType::get(llvmI8Ty);
-  Type llvmI32Ty = IntegerType::get(context, 32);
-  Type llvmI64Ty = IntegerType::get(context, 64);
-  Type llvmZTensorDescTy =
-      LLVM::LLVMPointerType::get(getZTensorDescStructTy(context));
-
-  // Got runtime error if using i64 as index to access zTensor. It looks
-  // like an error in MLIR. So use i32 here, which does not affect the
-  // correctness of the generated program.
-  Value zero = create.llvm.constant(llvmI32Ty, (int64_t)0);
-  Value one = create.llvm.constant(llvmI32Ty, (int64_t)1);
-  Value two = create.llvm.constant(llvmI32Ty, (int64_t)2);
-  Value three = create.llvm.constant(llvmI32Ty, (int64_t)3);
-  Value four = create.llvm.constant(llvmI32Ty, (int64_t)4);
+  Type llvmZTensorTy = getZTensorStructTy(context);
+  Type llvmZTensorPtrTy = krnl::getPointerType(context, llvmZTensorTy);
 
   // 1. Set pre-transformed descriptor.
   Value zTensorPreTransformedDescPtr = create.llvm.getElemPtr(
-      LLVM::LLVMPointerType::get(llvmZTensorDescTy), zTensor, {zero, zero});
+      llvmZTensorPtrTy, llvmZTensorTy, zTensor, ArrayRef<LLVM::GEPArg>{0, 0});
   create.llvm.store(preTransformedDescPtr, zTensorPreTransformedDescPtr);
 
   // 2. Set transformed descriptor.
   Value zTensorTransformedDescPtr = create.llvm.getElemPtr(
-      LLVM::LLVMPointerType::get(llvmZTensorDescTy), zTensor, {zero, one});
+      llvmZTensorPtrTy, llvmZTensorTy, zTensor, ArrayRef<LLVM::GEPArg>{0, 1});
   create.llvm.store(transformedDescPtr, zTensorTransformedDescPtr);
 
   // 3. Set buffer_size.
   Value bufferSizePtr = create.llvm.getElemPtr(
-      LLVM::LLVMPointerType::get(llvmI64Ty), zTensor, {zero, two});
+      llvmZTensorPtrTy, llvmZTensorTy, zTensor, ArrayRef<LLVM::GEPArg>{0, 2});
   create.llvm.store(bufferSize, bufferSizePtr);
 
   // 4. Set buffer. Buffer was allocated in advance by the stickified memref.
   // So get the pointer from the stickified memref and set it to the zTensor.
   Value bufferPtr = create.llvm.getElemPtr(
-      LLVM::LLVMPointerType::get(llvmI8PtrTy), zTensor, {zero, three});
+      llvmZTensorPtrTy, llvmZTensorTy, zTensor, ArrayRef<LLVM::GEPArg>{0, 3});
   create.llvm.store(alignedBuffer, bufferPtr);
 
   // 5. Set is_transformed.
   Value isTransformedVal =
       create.llvm.constant(llvmI1Ty, (int64_t)((isTransformed) ? 1 : 0));
   Value isTransformedDescPtr = create.llvm.getElemPtr(
-      LLVM::LLVMPointerType::get(llvmI1Ty), zTensor, {zero, four});
+      llvmZTensorPtrTy, llvmZTensorTy, zTensor, ArrayRef<LLVM::GEPArg>{0, 4});
   create.llvm.store(isTransformedVal, isTransformedDescPtr);
 
   // 6. Set reserved (not currently used), not touch

src/Accelerators/NNPA/Conversion/ZLowToLLVM/ZLowToLLVMCommon.cpp

@@ -188,7 +188,8 @@ void addKrnlToLLVMPasses(
   pm.addNestedPass<func::FuncOp>(mlir::createConvertSCFToCFPass());
 
   pm.addPass(mlir::memref::createFoldMemRefAliasOpsPass());
-  pm.addPass(krnl::createConvertKrnlToLLVMPass(verifyInputTensors));
+  pm.addPass(krnl::createConvertKrnlToLLVMPass(
+      verifyInputTensors, /*useOpaquePointers=*/true));
   pm.addPass(mlir::createReconcileUnrealizedCastsPass());
   pm.addPass(mlir::createCanonicalizerPass());
 }

src/Compiler/CompilerPasses.cpp

@@ -63,6 +63,8 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace krnl {
 
+bool LLVM_USE_OPAQUE_POINTER = true;
+
 uint64_t KRNL_ENTRY_POINT_ID = 0;
 
 // Return true if the value owns the storge. A value defined by memref.alloc
@@ -275,9 +277,9 @@ void genSignatureFunction(ModuleOp &module,
   Type i8Type = IntegerType::get(context, 8);
   Type i32Type = IntegerType::get(context, 32);
   Type i64Type = IntegerType::get(context, 64);
-  Type i64PtrTy = LLVM::LLVMPointerType::get(i64Type);
-  Type i8PtrTy = LLVM::LLVMPointerType::get(i8Type);
-  Type i8PtrPtrTy = LLVM::LLVMPointerType::get(i8PtrTy);
+  Type i64PtrTy = getPointerType(context, i64Type);
+  Type i8PtrTy = getPointerType(context, i8Type);
+  Type i8PtrPtrTy = getPointerType(context, i8PtrTy);
 
   uint64_t numOfEntryPoints = entryGlobalOps.size();
 
@@ -300,16 +302,13 @@ void genSignatureFunction(ModuleOp &module,
     uint32_t index = 0;
     Value lastValue = array;
     for (const LLVM::GlobalOp &globalOp : entryGlobalOps) {
-      Value address = create.llvm.addressOf(globalOp);
-      Value zeroI64 = create.llvm.constant(i64Type, (int64_t)0);
-      Value strAddr =
-          create.llvm.getElemPtr(i8PtrTy, address, {zeroI64, zeroI64});
+      Value strAddr = krnl::getPtrToGlobalString(globalOp, loc, b);
       lastValue =
           create.llvm.insertValue(arrayType, lastValue, strAddr, {index++});
     }
 
     // The last element of the array is NULL.
-    Value nullPtr = create.llvm.nullI8Ptr();
+    Value nullPtr = create.llvm.null(getI8PointerType(context));
     lastValue =
         create.llvm.insertValue(arrayType, lastValue, nullPtr, {index++});
     create.llvm._return(lastValue);
@@ -339,16 +338,15 @@ void genSignatureFunction(ModuleOp &module,
               LLVM::ICmpPredicate::ne, numOfEntryPoints, nullPtr);
         }, /*then=*/
         [&](LLVMBuilder &createLLVM) {
-          Value zero = createLLVM.constant(i64Type, (int64_t)0);
-          Value numOfEntryPointsPtr =
-              createLLVM.getElemPtr(i64PtrTy, numOfEntryPoints, {zero});
+          Value numOfEntryPointsPtr = createLLVM.getElemPtr(
+              i64PtrTy, i64Type, numOfEntryPoints, ArrayRef<LLVM::GEPArg>{0});
           Value noep =
               createLLVM.constant(i64Type, (int64_t)entryGlobalOps.size());
           createLLVM.store(noep, numOfEntryPointsPtr);
         });
     // Emit code to return the entry point array.
     Value entryAddr = create.llvm.addressOf(entryArrayOp);
-    Value entryI8Ptr = create.llvm.bitcastI8PtrPtr(entryAddr);
+    Value entryI8Ptr = create.llvm.bitcast(i8PtrPtrTy, entryAddr);
     create.llvm._return(entryI8Ptr);
   }
 
@@ -388,10 +386,8 @@ void genSignatureFunction(ModuleOp &module,
       create.llvm.ifThenElse(/*cond=*/
           [&](LLVMBuilder &createLLVM) {
             // Read an entry point name.
-            Value address = createLLVM.addressOf(globalEntryPoint);
-            Value zeroI64 = createLLVM.constant(i64Type, (int64_t)0);
             Value entryI8Ptr =
-                createLLVM.getElemPtr(i8PtrTy, address, {zeroI64, zeroI64});
+                krnl::getPtrToGlobalString(globalEntryPoint, loc, b);
             // Compare it with the user's entry point name.
             FlatSymbolRefAttr StrncmpRef = krnl::getOrInsertStrncmp(b, module);
             Value length = createLLVM.constant(
@@ -404,13 +400,13 @@ void genSignatureFunction(ModuleOp &module,
           }, /*then=*/
           [&](LLVMBuilder &createLLVM) {
             Value sigAddr = createLLVM.addressOf(globalSignature);
-            Value sigI8Ptr = createLLVM.bitcastI8Ptr(sigAddr);
+            Value sigI8Ptr = createLLVM.bitcast(i8PtrTy, sigAddr);
             createLLVM._return(sigI8Ptr);
           });
     }
 
     // Return NULL if not found.
-    create.llvm._return(create.llvm.nullI8Ptr());
+    create.llvm._return(create.llvm.null(getI8PointerType(context)));
   }
 }
 
@@ -427,8 +423,9 @@ struct ConvertKrnlToLLVMPass
   ConvertKrnlToLLVMPass() = default;
   ConvertKrnlToLLVMPass(const ConvertKrnlToLLVMPass &pass)
       : PassWrapper<ConvertKrnlToLLVMPass, OperationPass<ModuleOp>>() {}
-  ConvertKrnlToLLVMPass(bool verifyInputTensors) {
+  ConvertKrnlToLLVMPass(bool verifyInputTensors, bool useOpaquePointers) {
     this->verifyInputTensors = verifyInputTensors;
+    this->useOpaquePointers = useOpaquePointers;
   }
 
   StringRef getArgument() const override { return "convert-krnl-to-llvm"; }
@@ -439,6 +436,11 @@ struct ConvertKrnlToLLVMPass
 
   void runOnOperation() final;
 
+  Option<bool> useOpaquePointers{*this, "use-opaque-pointers",
+      llvm::cl::desc("Whether to use opaque pointers instead of typed pointers "
+                     "when lowering to LLVM. Default: true"),
+      llvm::cl::init(true)};
+
   Option<bool> verifyInputTensors{*this, "verify-input-tensors",
       llvm::cl::desc(
           "Verify input tensors whenever the entry point function is called.\n"
@@ -453,9 +455,10 @@ void ConvertKrnlToLLVMPass::runOnOperation() {
   const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
   LowerToLLVMOptions options(ctx, dataLayoutAnalysis.getAtOrAbove(module));
 
-  // There are many places where we still rely on non-opaque pointers. Disable
-  // opaque-pointers until we migrated the affected code parts
-  options.useOpaquePointers = false;
+  // MLIR/LLVM is moving to using opaque pointers instead of typed pointers.
+  // Remove this once MLIR/LLVM completely uses opaque pointers.
+  options.useOpaquePointers = useOpaquePointers; // for LLVMTypeConverter.
+  LLVM_USE_OPAQUE_POINTER = useOpaquePointers; // for onnx-mlir util functions.
 
   KRNL_ENTRY_POINT_ID = 0;
 
@@ -533,8 +536,10 @@ void ConvertKrnlToLLVMPass::runOnOperation() {
 std::unique_ptr<Pass> createConvertKrnlToLLVMPass() {
   return std::make_unique<ConvertKrnlToLLVMPass>();
 }
-std::unique_ptr<Pass> createConvertKrnlToLLVMPass(bool verifyInputTensors) {
-  return std::make_unique<ConvertKrnlToLLVMPass>(verifyInputTensors);
+std::unique_ptr<Pass> createConvertKrnlToLLVMPass(
+    bool verifyInputTensors, bool useOpaquePointers) {
+  return std::make_unique<ConvertKrnlToLLVMPass>(
+      verifyInputTensors, useOpaquePointers);
 }
 
 void populateKrnlToLLVMConversion(LLVMTypeConverter &typeConverter,