[CUDA] Add workaround for cuda driver bug (#6780)
The bug https://bugs.llvm.org/show_bug.cgi?id=48771 happens easily from
the code generated by scf.for lowering. To workaround it we mark all the
loops as nounroll right before llvm backend. This can be removed once we
move to an older driver for all the system we care about.
This also enables model tests that are now working.
diff --git a/iree/compiler/Dialect/HAL/Target/CUDA/BUILD b/iree/compiler/Dialect/HAL/Target/CUDA/BUILD
index 282aea8..41236c5 100644
--- a/iree/compiler/Dialect/HAL/Target/CUDA/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/CUDA/BUILD
@@ -33,9 +33,11 @@
name = "CUDA",
srcs = [
"CUDATarget.cpp",
+ "NoLoopUnrollPass.cpp",
],
hdrs = [
"CUDATarget.h",
+ "LLVMPasses.h",
],
deps = [
":cuda_libdevice",
@@ -45,6 +47,7 @@
"//iree/compiler/Dialect/HAL/Target",
"//iree/compiler/Utils",
"//iree/schemas:cuda_executable_def_c_fbs",
+ "@llvm-project//llvm:Analysis",
"@llvm-project//llvm:BitReader",
"@llvm-project//llvm:Core",
"@llvm-project//llvm:IPO",
diff --git a/iree/compiler/Dialect/HAL/Target/CUDA/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/CUDA/CMakeLists.txt
index 7ebad2b..c290659 100644
--- a/iree/compiler/Dialect/HAL/Target/CUDA/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/CUDA/CMakeLists.txt
@@ -33,10 +33,13 @@
CUDA
HDRS
"CUDATarget.h"
+ "LLVMPasses.h"
SRCS
"CUDATarget.cpp"
+ "NoLoopUnrollPass.cpp"
DEPS
::cuda_libdevice
+ LLVMAnalysis
LLVMBitReader
LLVMCore
LLVMipo
diff --git a/iree/compiler/Dialect/HAL/Target/CUDA/CUDATarget.cpp b/iree/compiler/Dialect/HAL/Target/CUDA/CUDATarget.cpp
index 2c0a276..1ce428e 100644
--- a/iree/compiler/Dialect/HAL/Target/CUDA/CUDATarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/CUDA/CUDATarget.cpp
@@ -7,6 +7,7 @@
#include "iree/compiler/Dialect/HAL/Target/CUDA/CUDATarget.h"
#include "iree/compiler/Codegen/Passes.h"
+#include "iree/compiler/Dialect/HAL/Target/CUDA/LLVMPasses.h"
#include "iree/compiler/Dialect/HAL/Target/CUDA/libdevice.h"
#include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
#include "iree/compiler/Utils/FlatbufferUtils.h"
@@ -44,6 +45,12 @@
llvm::raw_string_ostream stream(targetISA);
llvm::buffer_ostream pstream(stream);
llvm::legacy::PassManager codegenPasses;
+ // Workaround for CUDA driver bug
+ // (https://bugs.llvm.org/show_bug.cgi?id=48771), we mark all the loops with
+ // the no unroll metadata. This bug is fixed in cuda 11.4 but since we still
+ // run on older driver we need to keep it.
+ // TODO(thomasraoux): Remove it once we stop supporting older drivers.
+ codegenPasses.add(llvm::createSetNoUnrollPass());
targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
llvm::CGFT_AssemblyFile);
codegenPasses.run(module);
diff --git a/iree/compiler/Dialect/HAL/Target/CUDA/LLVMPasses.h b/iree/compiler/Dialect/HAL/Target/CUDA/LLVMPasses.h
new file mode 100644
index 0000000..cf997b1
--- /dev/null
+++ b/iree/compiler/Dialect/HAL/Target/CUDA/LLVMPasses.h
@@ -0,0 +1,19 @@
+// Copyright 2021 The IREE Authors
+//
+// 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
+
+#ifndef IREE_COMPILER_DIALECT_HAL_TARGET_CUDA_PASS_H_
+#define IREE_COMPILER_DIALECT_HAL_TARGET_CUDA_PASS_H_
+
+namespace llvm {
+
+class Pass;
+
+/// Pass to mark all loops with llvm metadata to disable unrolling.
+Pass *createSetNoUnrollPass();
+
+} // namespace llvm
+
+#endif // IREE_COMPILER_DIALECT_HAL_TARGET_CUDA_PASS_H_
diff --git a/iree/compiler/Dialect/HAL/Target/CUDA/NoLoopUnrollPass.cpp b/iree/compiler/Dialect/HAL/Target/CUDA/NoLoopUnrollPass.cpp
new file mode 100644
index 0000000..dfb6059
--- /dev/null
+++ b/iree/compiler/Dialect/HAL/Target/CUDA/NoLoopUnrollPass.cpp
@@ -0,0 +1,45 @@
+// Copyright 2021 The IREE Authors
+//
+// 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 "iree/compiler/Dialect/HAL/Target/CUDA/LLVMPasses.h"
+#include "llvm/Analysis/LoopPass.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "iree-dialect-hal-cuda-llvm-nounroll"
+
+namespace llvm {
+void initializeNoLoopUnrollPass(PassRegistry &Registry);
+}
+
+namespace {
+/// Pass that mark all loops with llvm.loop.unroll.disable metadata.
+class NoLoopUnroll : public LoopPass {
+ public:
+ static char ID;
+ NoLoopUnroll() : LoopPass(ID) {
+ initializeNoLoopUnrollPass(*PassRegistry::getPassRegistry());
+ }
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ AU.setPreservesCFG();
+ }
+ bool runOnLoop(Loop *L, LPPassManager &LPM) override {
+ L->setLoopAlreadyUnrolled();
+ return true;
+ }
+ StringRef getPassName() const override { return "Set Nounroll pass"; }
+};
+
+} // namespace
+
+char NoLoopUnroll::ID = 0;
+
+INITIALIZE_PASS_BEGIN(NoLoopUnroll, DEBUG_TYPE, "Set Nounroll", false, false)
+INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+INITIALIZE_PASS_END(NoLoopUnroll, DEBUG_TYPE, "Set Nounroll", false, false)
+
+Pass *llvm::createSetNoUnrollPass() { return new NoLoopUnroll(); }
diff --git a/iree/test/e2e/models/BUILD b/iree/test/e2e/models/BUILD
index 976a358..1732ac0 100644
--- a/iree/test/e2e/models/BUILD
+++ b/iree/test/e2e/models/BUILD
@@ -72,3 +72,20 @@
driver = "vulkan",
target_backend = "vulkan-spirv",
)
+
+iree_check_single_backend_test_suite(
+ name = "check_linalg_on_tensors_cuda_cuda",
+ timeout = "long",
+ srcs = CHECK_FRAMEWORK_TESTS,
+ compiler_flags = ["-iree-input-type=mhlo"],
+ driver = "cuda",
+ tags = [
+ # CUDA cuInit fails with sanitizer on.
+ "noasan",
+ "nomsan",
+ "notsan",
+ "noubsan",
+ "requires-gpu-nvidia",
+ ],
+ target_backend = "cuda",
+)
diff --git a/iree/test/e2e/models/CMakeLists.txt b/iree/test/e2e/models/CMakeLists.txt
index e8d3983..3975477 100644
--- a/iree/test/e2e/models/CMakeLists.txt
+++ b/iree/test/e2e/models/CMakeLists.txt
@@ -57,4 +57,24 @@
"-iree-input-type=mhlo"
)
+iree_check_single_backend_test_suite(
+ NAME
+ check_linalg_on_tensors_cuda_cuda
+ SRCS
+ "bert_encoder_unrolled_fake_weights.mlir"
+ "mobilenetv3_fake_weights.mlir"
+ TARGET_BACKEND
+ "cuda"
+ DRIVER
+ "cuda"
+ COMPILER_FLAGS
+ "-iree-input-type=mhlo"
+ LABELS
+ "noasan"
+ "nomsan"
+ "notsan"
+ "noubsan"
+ "requires-gpu-nvidia"
+)
+
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
