diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h --- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h @@ -123,6 +123,12 @@ /// annotation. void registerGpuSerializeToHsacoPass(); +/// Create an instance of the GPU kernel function to CUBIN binary serialization +/// pass. +std::unique_ptr createGpuSerializeToCubinPass(StringRef triple, + StringRef chip, + StringRef features); + /// Create an instance of the GPU kernel function to HSAco binary serialization /// pass. std::unique_ptr createGpuSerializeToHsacoPass(StringRef triple, diff --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp --- a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp @@ -48,7 +48,8 @@ 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 @@ } // namespace // Sets the 'option' to 'value' unless it already has a value. -static void maybeSetOption(Pass::Option &option, - const char *value) { +static void maybeSetOption(Pass::Option &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 @@ return std::make_unique(); }); } + +std::unique_ptr mlir::createGpuSerializeToCubinPass(StringRef triple, + StringRef arch, + StringRef features) { + return std::make_unique(triple, arch, features); +} + #else // MLIR_GPU_TO_CUBIN_PASS_ENABLE void mlir::registerGpuSerializeToCubinPass() {} #endif // MLIR_GPU_TO_CUBIN_PASS_ENABLE