[Mlir-commits] [mlir] 2f7a774 - [MLIR] Add a create function for mlir::SerializeToCubinPass

[William S. Moses]

Author: Ivan Radanov Ivanov
Date: 2022-09-21T18:02:59-04:00
New Revision: 2f7a774ed77ad549f4fc46ccf3637686184df240

URL: https://github.com/llvm/llvm-project/commit/2f7a774ed77ad549f4fc46ccf3637686184df240
DIFF: https://github.com/llvm/llvm-project/commit/2f7a774ed77ad549f4fc46ccf3637686184df240.diff

LOG: [MLIR] Add a create function for mlir::SerializeToCubinPass

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
    mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
index 09cbecfa5c10d..44429c3350cb5 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
@@ -123,6 +123,12 @@ void registerGpuSerializeToCubinPass();
 /// annotation.
 void registerGpuSerializeToHsacoPass();
 
+/// Create an instance of the GPU kernel function to CUBIN binary serialization
+/// pass.
+std::unique_ptr<Pass> createGpuSerializeToCubinPass(StringRef triple,
+                                                    StringRef chip,
+                                                    StringRef features);
+
 /// Create an instance of the GPU kernel function to HSAco binary serialization
 /// pass.
 std::unique_ptr<Pass> createGpuSerializeToHsacoPass(StringRef triple,

diff  --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
index e1b1e7e93d658..0f8c4488ebcf5 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
@@ -48,7 +48,8 @@ class SerializeToCubinPass
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SerializeToCubinPass)
 
-  SerializeToCubinPass();
+  SerializeToCubinPass(StringRef triple = "nvptx64-nvidia-cuda",
+                       StringRef chip = "sm_80", StringRef features = "+ptx75");
 
   StringRef getArgument() const override { return "gpu-to-cubin"; }
   StringRef getDescription() const override {
@@ -65,16 +66,16 @@ class SerializeToCubinPass
 } // namespace
 
 // Sets the 'option' to 'value' unless it already has a value.
-static void maybeSetOption(Pass::Option<std::string> &option,
-                           const char *value) {
+static void maybeSetOption(Pass::Option<std::string> &option, StringRef value) {
   if (!option.hasValue())
-    option = value;
+    option = value.str();
 }
 
-SerializeToCubinPass::SerializeToCubinPass() {
-  maybeSetOption(this->triple, "nvptx64-nvidia-cuda");
-  maybeSetOption(this->chip, "sm_35");
-  maybeSetOption(this->features, "+ptx60");
+SerializeToCubinPass::SerializeToCubinPass(StringRef triple, StringRef chip,
+                                           StringRef features) {
+  maybeSetOption(this->triple, triple);
+  maybeSetOption(this->chip, chip);
+  maybeSetOption(this->features, features);
 }
 
 void SerializeToCubinPass::getDependentDialects(
@@ -143,6 +144,13 @@ void mlir::registerGpuSerializeToCubinPass() {
     return std::make_unique<SerializeToCubinPass>();
   });
 }
+
+std::unique_ptr<Pass> mlir::createGpuSerializeToCubinPass(StringRef triple,
+                                                          StringRef arch,
+                                                          StringRef features) {
+  return std::make_unique<SerializeToCubinPass>(triple, arch, features);
+}
+
 #else  // MLIR_GPU_TO_CUBIN_PASS_ENABLE
 void mlir::registerGpuSerializeToCubinPass() {}
 #endif // MLIR_GPU_TO_CUBIN_PASS_ENABLE