@@ -75,32 +75,27 @@ OffloadTargetInfo::OffloadTargetInfo(const StringRef Target,
7575 const OffloadBundlerConfig &BC)
7676 : BundlerConfig(BC) {
7777
78- // TODO: Add error checking from ClangOffloadBundler.cpp
79- auto TargetFeatures = Target.split (' :' );
80- auto TripleOrGPU = TargetFeatures.first .rsplit (' -' );
81-
82- if (clang::StringToOffloadArch (TripleOrGPU.second ) !=
83- clang::OffloadArch::UNKNOWN) {
84- auto KindTriple = TripleOrGPU.first .split (' -' );
85- this ->OffloadKind = KindTriple.first ;
86-
87- // Enforce optional env field to standardize bundles
88- llvm::Triple t = llvm::Triple (KindTriple.second );
89- this ->Triple = llvm::Triple (t.getArchName (), t.getVendorName (),
90- t.getOSName (), t.getEnvironmentName ());
91-
92- this ->TargetID = Target.substr (Target.find (TripleOrGPU.second ));
93- } else {
94- auto KindTriple = TargetFeatures.first .split (' -' );
95- this ->OffloadKind = KindTriple.first ;
96-
97- // Enforce optional env field to standardize bundles
98- llvm::Triple t = llvm::Triple (KindTriple.second );
99- this ->Triple = llvm::Triple (t.getArchName (), t.getVendorName (),
100- t.getOSName (), t.getEnvironmentName ());
101-
78+ // <kind>-<triple>[-<target id>[:target features]]
79+ // <triple> := <arch>-<vendor>-<os>-<env>
80+ SmallVector<StringRef, 6 > Components;
81+ Target.split (Components, ' -' , /* MaxSplit=*/ 5 );
82+ assert ((Components.size () == 5 || Components.size () == 6 ) &&
83+ " malformed target string" );
84+
85+ StringRef TargetIdWithFeature =
86+ Components.size () == 6 ? Components.back () : " " ;
87+ StringRef TargetId = TargetIdWithFeature.split (' :' ).first ;
88+ if (!TargetId.empty () &&
89+ clang::StringToOffloadArch (TargetId) != clang::OffloadArch::UNKNOWN)
90+ this ->TargetID = TargetIdWithFeature;
91+ else
10292 this ->TargetID = " " ;
103- }
93+
94+ this ->OffloadKind = Components.front ();
95+ ArrayRef<StringRef> TripleSlice{&Components[1 ], /* length=*/ 4 };
96+ llvm::Triple T = llvm::Triple (llvm::join (TripleSlice, " -" ));
97+ this ->Triple = llvm::Triple (T.getArchName (), T.getVendorName (), T.getOSName (),
98+ T.getEnvironmentName ());
10499}
105100
106101bool OffloadTargetInfo::hasHostKind () const {
@@ -140,7 +135,18 @@ bool OffloadTargetInfo::operator==(const OffloadTargetInfo &Target) const {
140135}
141136
142137std::string OffloadTargetInfo::str () const {
143- return Twine (OffloadKind + " -" + Triple.str () + " -" + TargetID).str ();
138+ std::string NormalizedTriple;
139+ // Unfortunately we need some special sauce for AMDGPU because all the runtime
140+ // assumes the triple to be "amdgcn-amd-amdhsa-" (empty environment) instead
141+ // of "amdgcn-amd-amdhsa-unknown". It's gonna be very tricky to patch
142+ // different layers of runtime.
143+ if (Triple.isAMDGPU ()) {
144+ NormalizedTriple = Triple.normalize (Triple::CanonicalForm::THREE_IDENT);
145+ NormalizedTriple.push_back (' -' );
146+ } else {
147+ NormalizedTriple = Triple.normalize (Triple::CanonicalForm::FOUR_IDENT);
148+ }
149+ return Twine (OffloadKind + " -" + NormalizedTriple + " -" + TargetID).str ();
144150}
145151
146152static StringRef getDeviceFileExtension (StringRef Device,
@@ -1424,6 +1430,9 @@ Error OffloadBundler::UnbundleFiles() {
14241430 StringMap<StringRef> Worklist;
14251431 auto Output = BundlerConfig.OutputFileNames .begin ();
14261432 for (auto &Triple : BundlerConfig.TargetNames ) {
1433+ if (!checkOffloadBundleID (Triple))
1434+ return createStringError (errc::invalid_argument,
1435+ " invalid bundle id from bundle config" );
14271436 Worklist[Triple] = *Output;
14281437 ++Output;
14291438 }
@@ -1443,6 +1452,9 @@ Error OffloadBundler::UnbundleFiles() {
14431452
14441453 StringRef CurTriple = **CurTripleOrErr;
14451454 assert (!CurTriple.empty ());
1455+ if (!checkOffloadBundleID (CurTriple))
1456+ return createStringError (errc::invalid_argument,
1457+ " invalid bundle id read from the bundle" );
14461458
14471459 auto Output = Worklist.begin ();
14481460 for (auto E = Worklist.end (); Output != E; Output++) {
@@ -1501,6 +1513,8 @@ Error OffloadBundler::UnbundleFiles() {
15011513 return createFileError (E.second , EC);
15021514
15031515 // If this entry has a host kind, copy the input file to the output file.
1516+ // We don't need to check E.getKey() here through checkOffloadBundleID
1517+ // because the entire WorkList has been checked above.
15041518 auto OffloadInfo = OffloadTargetInfo (E.getKey (), BundlerConfig);
15051519 if (OffloadInfo.hasHostKind ())
15061520 OutputFile.write (Input.getBufferStart (), Input.getBufferSize ());
@@ -1730,6 +1744,10 @@ Error OffloadBundler::UnbundleArchive() {
17301744 // archive.
17311745 while (!CodeObject.empty ()) {
17321746 SmallVector<StringRef> CompatibleTargets;
1747+ if (!checkOffloadBundleID (CodeObject)) {
1748+ return createStringError (errc::invalid_argument,
1749+ " Invalid bundle id read from code object" );
1750+ }
17331751 auto CodeObjectInfo = OffloadTargetInfo (CodeObject, BundlerConfig);
17341752 if (getCompatibleOffloadTargets (CodeObjectInfo, CompatibleTargets,
17351753 BundlerConfig)) {
@@ -1819,3 +1837,11 @@ Error OffloadBundler::UnbundleArchive() {
18191837
18201838 return Error::success ();
18211839}
1840+
1841+ bool clang::checkOffloadBundleID (const llvm::StringRef Str) {
1842+ // <kind>-<triple>[-<target id>[:target features]]
1843+ // <triple> := <arch>-<vendor>-<os>-<env>
1844+ SmallVector<StringRef, 6 > Components;
1845+ Str.split (Components, ' -' , /* MaxSplit=*/ 5 );
1846+ return Components.size () == 5 || Components.size () == 6 ;
1847+ }
0 commit comments