unittests/Target/LLVM/SerializeNVVMTarget.cpp

//===- SerializeNVVMTarget.cpp ----------------------------------*- C++ -*-===//
//
// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Config/mlir-config.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Remarks.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/InitAllDialects.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Target/LLVM/NVVM/Target.h"
#include "mlir/Target/LLVM/NVVM/Utils.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"

#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/Config/Targets.h" // for LLVM_HAS_NVPTX_TARGET
#include "llvm/IRReader/IRReader.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/LibraryOptionsContext.h"
#include "llvm/Support/MemoryBufferRef.h"
#include "llvm/Support/Process.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/TargetParser/Host.h"

#include "gmock/gmock.h"
#include <cstdint>

using namespace mlir;

// Skip the test if the NVPTX target was not built.
#if LLVM_HAS_NVPTX_TARGET
#define SKIP_WITHOUT_NVPTX(x) x
#else
#define SKIP_WITHOUT_NVPTX(x) DISABLED_##x
#endif

class MLIRTargetLLVMNVVM : public ::testing::Test {
protected:
  void SetUp() override {
    registerBuiltinDialectTranslation(registry);
    registerLLVMDialectTranslation(registry);
    registerGPUDialectTranslation(registry);
    registerNVVMDialectTranslation(registry);
    NVVM::registerNVVMTargetInterfaceExternalModels(registry);
  }

  // Checks if PTXAS is in PATH.
  bool hasPtxas() {
    // Find the `ptxas` compiler.
    std::optional<std::string> ptxasCompiler =
        llvm::sys::Process::FindInEnvPath("PATH", "ptxas");
    return ptxasCompiler.has_value();
  }

  // Dialect registry.
  DialectRegistry registry;

  // MLIR module used for the tests.
  const std::string moduleStr = R"mlir(
      gpu.module @nvvm_test {
        llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {
        llvm.return
      }
    })mlir";
};

// Test NVVM serialization to LLVM.
TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMMToLLVM)) {
  MLIRContext context(registry);

  OwningOpRef<ModuleOp> module =
      parseSourceString<ModuleOp>(moduleStr, &context);
  ASSERT_TRUE(!!module);

  // Create an NVVM target.
  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);

  // Serialize the module.
  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
  ASSERT_TRUE(!!serializer);
  gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Offload);
  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
    std::optional<mlir::gpu::SerializedObject> object =
        serializer.serializeToObject(gpuModule, options);
    // Check that the serializer was successful.
    ASSERT_TRUE(object != std::nullopt);
    ASSERT_TRUE(!object->getObject().empty());

    // Read the serialized module.
    llvm::MemoryBufferRef buffer(
        StringRef(object->getObject().data(), object->getObject().size()),
        "module");
    llvm::LLVMContext llvmContext;
    llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =
        llvm::getLazyBitcodeModule(buffer, llvmContext);
    ASSERT_TRUE(!!llvmModule);
    ASSERT_TRUE(!!*llvmModule);

    // Check that it has a function named `foo`.
    ASSERT_TRUE((*llvmModule)->getFunction("nvvm_kernel") != nullptr);
  }
}

// Test NVVM serialization to PTX.
TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {
  MLIRContext context(registry);

  OwningOpRef<ModuleOp> module =
      parseSourceString<ModuleOp>(moduleStr, &context);
  ASSERT_TRUE(!!module);

  // Create an NVVM target.
  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);

  // Serialize the module.
  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
  ASSERT_TRUE(!!serializer);
  gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
    std::optional<mlir::gpu::SerializedObject> object =
        serializer.serializeToObject(gpuModule, options);
    // Check that the serializer was successful.
    ASSERT_TRUE(object != std::nullopt);
    ASSERT_TRUE(!object->getObject().empty());

    ASSERT_TRUE(
        StringRef(object->getObject().data(), object->getObject().size())
            .contains("nvvm_kernel"));
    ASSERT_TRUE(
        StringRef(object->getObject().data(), object->getObject().size())
            .count('\0') == 0);
  }
}

// The LLVMContext and TargetMachine that serialization creates read the library
// options in TargetOptions.
TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(LibraryOptions)) {
  MLIRContext context(registry);
  context.loadDialect<NVVM::NVVMDialect>();
  OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(R"mlir(
      gpu.module @nvvm_test {
        llvm.func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
          %0 = llvm.fmul %a, %b : f32
          %1 = llvm.fadd %0, %c : f32
          llvm.return %1 : f32
        }
      })mlir",
                                                             &context);
  ASSERT_TRUE(!!module);

  // Register the NVPTX options before parsing them.
  NVVM::SerializeGPUModuleBase::init();
  llvm::cl::LibraryOptionsContext libraryOptions;
  const char *args[] = {"test", "-nvptx-fma-level=1"};
  ASSERT_TRUE(llvm::cl::parseCommandLineOptions(libraryOptions, args, "",
                                                &llvm::nulls()));

  auto serializer =
      dyn_cast<gpu::TargetAttrInterface>(NVVM::NVVMTargetAttr::get(&context));
  ASSERT_TRUE(!!serializer);
  size_t numArgs = 0;
  auto initialCallback = [&](llvm::Module &m) {
    numArgs = m.getContext().getLibraryOptions().args().size();
  };
  gpu::TargetOptions options({}, {}, {}, {}, gpu::CompilationTarget::Assembly,
                             {}, initialCallback);
  options.setLibraryOptions(&libraryOptions);
  std::optional<gpu::SerializedObject> object = serializer.serializeToObject(
      *(*module).getBody()->getOps<gpu::GPUModuleOp>().begin(), options);
  ASSERT_TRUE(object != std::nullopt);
  EXPECT_THAT(StringRef(object->getObject().data(), object->getObject().size()),
              testing::HasSubstr("fma.rn.f32"));
  EXPECT_EQ(numArgs, 1u);
}

// Test NVVM serialization to Binary.
TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {
  if (!hasPtxas())
    GTEST_SKIP() << "PTXAS compiler not found, skipping test.";

  MLIRContext context(registry);

  OwningOpRef<ModuleOp> module =
      parseSourceString<ModuleOp>(moduleStr, &context);
  ASSERT_TRUE(!!module);

  // Create an NVVM target.
  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);

  // Serialize the module.
  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
  ASSERT_TRUE(!!serializer);
  gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);
  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
    std::optional<mlir::gpu::SerializedObject> object =
        serializer.serializeToObject(gpuModule, options);
    // Check that the serializer was successful.
    ASSERT_TRUE(object != std::nullopt);
    ASSERT_TRUE(!object->getObject().empty());
  }
}

// Test callback functions invoked with LLVM IR and ISA.
TEST_F(MLIRTargetLLVMNVVM,
       SKIP_WITHOUT_NVPTX(CallbackInvokedWithLLVMIRAndISA)) {
  MLIRContext context(registry);

  OwningOpRef<ModuleOp> module =
      parseSourceString<ModuleOp>(moduleStr, &context);
  ASSERT_TRUE(!!module);

  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);

  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
  ASSERT_TRUE(!!serializer);

  std::string initialLLVMIR;
  auto initialCallback = [&initialLLVMIR](llvm::Module &module) {
    llvm::raw_string_ostream ros(initialLLVMIR);
    module.print(ros, nullptr);
  };

  std::string linkedLLVMIR;
  auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
    llvm::raw_string_ostream ros(linkedLLVMIR);
    module.print(ros, nullptr);
  };

  std::string optimizedLLVMIR;
  auto optimizedCallback = [&optimizedLLVMIR](llvm::Module &module) {
    llvm::raw_string_ostream ros(optimizedLLVMIR);
    module.print(ros, nullptr);
  };

  std::string isaResult;
  auto isaCallback = [&isaResult](llvm::StringRef isa) {
    isaResult = isa.str();
  };

  gpu::TargetOptions options({}, {}, {}, {}, gpu::CompilationTarget::Assembly,
                             {}, initialCallback, linkedCallback,
                             optimizedCallback, isaCallback);

  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
    std::optional<mlir::gpu::SerializedObject> object =
        serializer.serializeToObject(gpuModule, options);

    ASSERT_TRUE(object != std::nullopt);
    ASSERT_TRUE(!object->getObject().empty());
    ASSERT_TRUE(!initialLLVMIR.empty());
    ASSERT_TRUE(!linkedLLVMIR.empty());
    ASSERT_TRUE(!optimizedLLVMIR.empty());
    ASSERT_TRUE(!isaResult.empty());

    initialLLVMIR.clear();
    linkedLLVMIR.clear();
    optimizedLLVMIR.clear();
    isaResult.clear();
  }
}

// Test that LLVM optimization remarks are imported into the remark engine.
TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(LLVMRemarksAreImported)) {
  MLIRContext context(registry);
  std::string moduleStr = R"mlir(
    gpu.module @nvvm_test {
      llvm.func @helper(%a: f32, %b: f32) -> f32 {
        %0 = llvm.fadd %a, %b : f32
        llvm.return %0 : f32
      }
      llvm.func @nvvm_kernel(%arg0: f32, %arg1: !llvm.ptr)
          attributes {gpu.kernel, nvvm.kernel} {
        %0 = llvm.call @helper(%arg0, %arg0) : (f32, f32) -> f32
        llvm.store %0, %arg1 : f32, !llvm.ptr
        llvm.return
      }
    }
  )mlir";
  OwningOpRef<ModuleOp> module =
      parseSourceString<ModuleOp>(moduleStr, &context);
  ASSERT_TRUE(!!module);

  struct CollectingStreamer : public remark::detail::MLIRRemarkStreamerBase {
    CollectingStreamer(std::vector<remark::detail::Remark> &remarks)
        : remarks(remarks) {}
    void
    streamOptimizationRemark(const remark::detail::Remark &remark) override {
      remarks.push_back(remark);
    }
    std::vector<remark::detail::Remark> &remarks;
  };
  std::vector<remark::detail::Remark> remarks;
  remark::RemarkCategories cats{/*all=*/std::nullopt, /*passed=*/"llvm-inline",
                                /*missed=*/std::nullopt,
                                /*analysis=*/"llvm-asm-printer",
                                /*failed=*/std::nullopt};
  ASSERT_TRUE(succeeded(remark::enableOptimizationRemarks(
      context, std::make_unique<CollectingStreamer>(remarks),
      std::make_unique<remark::RemarkEmittingPolicyAll>(), cats)));

  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
  ASSERT_TRUE(!!serializer);
  gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
    std::optional<mlir::gpu::SerializedObject> object =
        serializer.serializeToObject(gpuModule, options);
    ASSERT_TRUE(object != std::nullopt);
    ASSERT_TRUE(!object->getObject().empty());

    Operation *kernel =
        SymbolTable::lookupSymbolIn(gpuModule.getOperation(), "nvvm_kernel");
    ASSERT_TRUE(kernel);

    // Middle-end remark.
    auto inlined =
        llvm::find_if(remarks, [](const remark::detail::Remark &remark) {
          return remark.getCategoryName() == "llvm-inline" &&
                 remark.getRemarkName() == "Inlined";
        });
    ASSERT_TRUE(inlined != remarks.end());
    EXPECT_EQ(inlined->getRemarkKind(), remark::RemarkKind::RemarkPassed);
    EXPECT_EQ(inlined->getFunction(), "nvvm_kernel");
    EXPECT_EQ(inlined->getLocation(), kernel->getLoc());

    // Codegen remark.
    auto instructionCount =
        llvm::find_if(remarks, [](const remark::detail::Remark &remark) {
          return remark.getCategoryName() == "llvm-asm-printer" &&
                 remark.getRemarkName() == "InstructionCount" &&
                 remark.getFunction() == "nvvm_kernel";
        });
    ASSERT_TRUE(instructionCount != remarks.end());
    EXPECT_EQ(instructionCount->getRemarkKind(),
              remark::RemarkKind::RemarkAnalysis);
    EXPECT_EQ(instructionCount->getLocation(), kernel->getLoc());
    remarks.clear();
  }
}

// Test linking LLVM IR from a resource attribute.
TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(LinkedLLVMIRResource)) {
  MLIRContext context(registry);
  std::string moduleStr = R"mlir(
    gpu.module @nvvm_test {
      llvm.func @bar()
      llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {
        llvm.call @bar() : () -> ()
        llvm.return
      }
    }
  )mlir";
  // Provide the library to link as a serialized bitcode blob.
  SmallVector<char> bitcodeToLink;
  {
    std::string linkedLib = R"llvm(
  define void @bar() {
    ret void
  }
  )llvm";
    llvm::SMDiagnostic err;
    llvm::MemoryBufferRef buffer(linkedLib, "linkedLib");
    llvm::LLVMContext llvmCtx;
    std::unique_ptr<llvm::Module> module = llvm::parseIR(buffer, err, llvmCtx);
    ASSERT_TRUE(module) << " Can't parse IR: " << err.getMessage();
    {
      llvm::raw_svector_ostream os(bitcodeToLink);
      WriteBitcodeToFile(*module, os);
    }
  }

  OwningOpRef<ModuleOp> module =
      parseSourceString<ModuleOp>(moduleStr, &context);
  ASSERT_TRUE(!!module);
  Builder builder(&context);

  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);

  // Hook to intercept the LLVM IR after linking external libs.
  std::string linkedLLVMIR;
  auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
    llvm::raw_string_ostream ros(linkedLLVMIR);
    module.print(ros, nullptr);
  };

  // Store the bitcode as a DenseI8ArrayAttr.
  SmallVector<Attribute> librariesToLink;
  librariesToLink.push_back(DenseI8ArrayAttr::get(
      &context,
      ArrayRef<int8_t>((int8_t *)bitcodeToLink.data(), bitcodeToLink.size())));
  gpu::TargetOptions options({}, librariesToLink, {}, {},
                             gpu::CompilationTarget::Assembly, {}, {},
                             linkedCallback);
  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
    std::optional<mlir::gpu::SerializedObject> object =
        serializer.serializeToObject(gpuModule, options);

    // Verify that we correctly linked in the library: the external call is
    // replaced by the definition.
    ASSERT_TRUE(!linkedLLVMIR.empty());
    {
      llvm::SMDiagnostic err;
      llvm::MemoryBufferRef buffer(linkedLLVMIR, "linkedLLVMIR");
      llvm::LLVMContext llvmCtx;
      std::unique_ptr<llvm::Module> module =
          llvm::parseIR(buffer, err, llvmCtx);
      ASSERT_TRUE(module) << " Can't parse linkedLLVMIR: " << err.getMessage()
                          << " IR: \n\b" << linkedLLVMIR;
      llvm::Function *bar = module->getFunction("bar");
      ASSERT_TRUE(bar);
      ASSERT_FALSE(bar->empty());
    }
    ASSERT_TRUE(object != std::nullopt);
    ASSERT_TRUE(!object->getObject().empty());
  }
}