yum-slop/modular_slang

write individual HLSL modules/libraries in slang

git clone https://git.yummers.dev/yum-slop/modular_slang

yumaccelerate abseil build30c5fa6

master
branchesmasterall refs →
20.7 KiB759 linesraw
1#include <array>
2#include <filesystem>
3#include <cctype>
4#include <fstream>
5#include <iomanip>
6#include <iterator>
7#include <iostream>
8#include <limits>
9#include <sstream>
10#include <string>
11#include <string_view>
12#include <unordered_set>
13#include <vector>
14#include <utility>
15
16#include "absl/status/status.h"
17#include "absl/status/statusor.h"
18
19#include <slang.h>
20#include <slang-com-ptr.h>
21
22namespace fs = std::filesystem;
23
24using ::slang::CompilerOptionEntry;
25using ::slang::CompilerOptionName;
26using ::slang::createGlobalSession;
27using ::slang::DeclReflection;
28using ::slang::FunctionReflection;
29using ::slang::IBlob;
30using ::slang::ICompileRequest;
31using ::slang::IGlobalSession;
32using ::slang::IModule;
33using ::slang::ISession;
34using ::slang::SessionDesc;
35using ::slang::TargetDesc;
36
37template <typename T>
38using ComPtr = ::Slang::ComPtr<T>;
39
40// Print any diagnostics carried by a Slang blob with optional context information.
41void printDiagnostics(const char* context, IBlob* diagnostics) {
42  if (!diagnostics) {
43    return;
44  }
45
46  std::size_t size = diagnostics->getBufferSize();
47  if (size == 0) {
48    return;
49  }
50
51  std::string_view text(static_cast<const char*>(diagnostics->getBufferPointer()), size);
52  if (!text.empty() && text.back() == '\0') {
53    text.remove_suffix(1);
54  }
55
56  if (text.empty()) {
57    return;
58  }
59
60  if (context && *context) {
61    std::cerr << context << " diagnostics:" << std::endl;
62  }
63
64  std::cerr.write(text.data(), text.size());
65  if (text.back() != '\n') {
66    std::cerr << std::endl;
67  }
68}
69
70// Helper to convert Slang API results into absl::Status values.
71absl::Status checkSlangResult(const char* context, SlangResult res, IBlob* diagnostics = nullptr) {
72  printDiagnostics(context, diagnostics);
73
74  if (SLANG_FAILED(res)) {
75    std::ostringstream message;
76    message << (context && *context ? context : "Slang call")
77    << " failed with SlangResult " << res
78    << " (0x" << std::hex << res << std::dec << ')';
79    return absl::InternalError(message.str());
80  }
81
82  return absl::OkStatus();
83}
84
85absl::Status writeTextFile(const fs::path& path, std::string_view contents) {
86  std::ofstream file(path, std::ios::binary);
87  if (!file) {
88    std::ostringstream msg;
89    msg << "Failed to open " << path << " for writing.";
90    return absl::InternalError(msg.str());
91  }
92
93  file.write(contents.data(), static_cast<std::streamsize>(contents.size()));
94  file.close();
95
96  if (!file) {
97    std::ostringstream msg;
98    msg << "Failed to write " << path;
99    return absl::InternalError(msg.str());
100  }
101
102  return absl::OkStatus();
103}
104
105void addCompilerOption(std::vector<CompilerOptionEntry>& options, CompilerOptionName name) {
106  CompilerOptionEntry entry = {};
107  entry.name = name;
108  entry.value.intValue0 = 1;
109  options.push_back(entry);
110}
111
112struct FunctionInfo {
113  std::string name;
114};
115
116struct IncludeGuardInfo {
117  bool present = false;
118  std::string macro;
119  std::string ifndefLine;
120  std::string defineLine;
121  std::string endifLine;
122};
123
124struct ModuleRequest {
125  fs::path modulePath;
126  std::string moduleName;
127  std::string searchPath;
128  fs::path outputPath;
129};
130
131std::string trim(std::string_view text) {
132  std::size_t start = 0;
133  std::size_t end = text.size();
134
135  while (start < end && std::isspace(static_cast<unsigned char>(text[start]))) {
136    ++start;
137  }
138
139  while (end > start && std::isspace(static_cast<unsigned char>(text[end - 1]))) {
140    --end;
141  }
142
143  return std::string(text.substr(start, end - start));
144}
145
146bool isTopLevelFunction(DeclReflection* functionDecl) {
147  if (!functionDecl) {
148    return false;
149  }
150
151  using Kind = DeclReflection::Kind;
152  for (DeclReflection* parent = functionDecl->getParent(); parent;
153      parent = parent->getParent()) {
154    switch (parent->getKind()) {
155    case Kind::Module:
156    case Kind::Namespace:
157      return true;
158    case Kind::Generic:
159      continue;
160    default:
161      return false;
162    }
163  }
164
165  return false;
166}
167
168
169std::unordered_set<std::string> findPublicFunctionNames(const fs::path& sourcePath) {
170  std::unordered_set<std::string> names;
171
172  std::ifstream input(sourcePath, std::ios::binary);
173  if (!input) {
174    return names;
175  }
176
177  std::string source((std::istreambuf_iterator<char>(input)), std::istreambuf_iterator<char>());
178  const std::size_t length = source.size();
179
180  std::size_t index = 0;
181  bool publicPending = false;
182  std::string candidate;
183  int templateDepth = 0;
184
185  while (index < length) {
186    char c = source[index];
187
188    if (c == '/' && index + 1 < length) {
189      char next = source[index + 1];
190      if (next == '/') {
191        index += 2;
192        while (index < length && source[index] != '\n') {
193          ++index;
194        }
195        continue;
196      }
197      if (next == '*') {
198        index += 2;
199        while (index + 1 < length && !(source[index] == '*' && source[index + 1] == '/')) {
200          ++index;
201        }
202        if (index + 1 < length) {
203          index += 2;
204        }
205        continue;
206      }
207    }
208
209    if (c == '"' || c == '\'') {
210      char quote = c;
211      ++index;
212      while (index < length) {
213        char current = source[index];
214        if (current == '\\') {
215          index += 2;
216          continue;
217        }
218        if (current == quote) {
219          ++index;
220          break;
221        }
222        ++index;
223      }
224      continue;
225    }
226
227    if (std::isalpha(static_cast<unsigned char>(c)) || c == '_') {
228      std::size_t startToken = index;
229      ++index;
230      while (index < length) {
231        char ch = source[index];
232        if (std::isalnum(static_cast<unsigned char>(ch)) || ch == '_') {
233          ++index;
234        } else {
235          break;
236        }
237      }
238      std::string token = source.substr(startToken, index - startToken);
239      if (token == "public") {
240        publicPending = true;
241        candidate.clear();
242        templateDepth = 0;
243      } else if (publicPending && templateDepth == 0) {
244        candidate = token;
245      }
246      continue;
247    }
248
249    if (publicPending) {
250      if (c == '<') {
251        ++templateDepth;
252        ++index;
253        continue;
254      }
255      if (c == '>') {
256        if (templateDepth > 0) {
257          --templateDepth;
258        }
259        ++index;
260        continue;
261      }
262      if (c == '(') {
263        if (!candidate.empty() && templateDepth == 0) {
264          names.insert(candidate);
265        }
266        publicPending = false;
267        candidate.clear();
268        templateDepth = 0;
269        ++index;
270        continue;
271      }
272      if (c == ';' || c == '{' || c == '}') {
273        publicPending = false;
274        candidate.clear();
275        templateDepth = 0;
276        ++index;
277        continue;
278      }
279    }
280
281    ++index;
282  }
283
284  return names;
285}
286
287
288
289// Recursively gather function declarations defined in the supplied Slang module.
290void collectFunctionInfos(
291    DeclReflection* decl,
292    const std::unordered_set<std::string>& publicFunctions,
293    std::vector<FunctionInfo>& functions,
294    std::unordered_set<std::string>& seenNames) {
295  if (!decl) {
296    return;
297  }
298
299  using Kind = DeclReflection::Kind;
300
301  switch (decl->getKind()) {
302  case Kind::Func:
303    if (auto* functionReflection = decl->asFunction()) {
304      if (const char* name = functionReflection->getName()) {
305        bool isPublic = publicFunctions.find(name) != publicFunctions.end();
306
307        if (*name && isPublic && seenNames.insert(name).second && isTopLevelFunction(decl)) {
308          std::cerr << "Discovered entry point: " << name << std::endl;
309          functions.push_back({name});
310        }
311      }
312    }
313    break;
314  case Kind::Generic:
315    if (auto* genericDecl = decl->asGeneric()) {
316      collectFunctionInfos(
317          genericDecl->getInnerDecl(),
318          publicFunctions,
319          functions,
320          seenNames);
321    }
322    break;
323  default:
324    break;
325  }
326
327  for (auto* child : decl->getChildren()) {
328    collectFunctionInfos(child, publicFunctions, functions, seenNames);
329  }
330}
331
332
333IncludeGuardInfo detectIncludeGuard(const fs::path& sourcePath) {
334  IncludeGuardInfo info;
335
336  std::ifstream input(sourcePath);
337  if (!input) {
338    return info;
339  }
340
341  std::vector<std::string> lines;
342  std::string line;
343  while (std::getline(input, line)) {
344    lines.push_back(line);
345  }
346
347  std::size_t ifndefIndex = std::numeric_limits<std::size_t>::max();
348  for (std::size_t i = 0; i < lines.size(); ++i) {
349    std::string trimmed = trim(lines[i]);
350    if (trimmed.rfind("#ifndef", 0) == 0) {
351      std::istringstream stream(trimmed);
352      std::string directive;
353      std::string macro;
354      stream >> directive >> macro;
355      if (!macro.empty()) {
356        info.macro = macro;
357        info.ifndefLine = lines[i];
358        ifndefIndex = i;
359      }
360      break;
361    }
362  }
363
364  if (info.macro.empty()) {
365    return info;
366  }
367
368  for (std::size_t i = ifndefIndex + 1; i < lines.size(); ++i) {
369    std::string trimmed = trim(lines[i]);
370    if (trimmed.rfind("#define", 0) == 0) {
371      std::istringstream stream(trimmed);
372      std::string directive;
373      std::string macro;
374      stream >> directive >> macro;
375      if (macro == info.macro) {
376        info.defineLine = lines[i];
377        break;
378      }
379    }
380  }
381
382  if (info.defineLine.empty()) {
383    info = IncludeGuardInfo{};
384    return info;
385  }
386
387  for (std::size_t i = lines.size(); i-- > 0;) {
388    std::string trimmed = trim(lines[i]);
389    if (trimmed.rfind("#endif", 0) == 0) {
390      info.endifLine = lines[i];
391      break;
392    }
393  }
394
395  if (info.endifLine.empty()) {
396    info = IncludeGuardInfo{};
397    return info;
398  }
399
400  info.present = true;
401  return info;
402}
403
404absl::StatusOr<ModuleRequest> parseModuleRequest(int argc, char** argv) {
405  const char* programName = (argc > 0 && argv) ? argv[0] : "modular_slang";
406
407  if (argc < 2 || !argv) {
408    std::ostringstream usage;
409    usage << "Usage: " << programName << " <module.slang>";
410    return absl::InvalidArgumentError(usage.str());
411  }
412
413  ModuleRequest request;
414  request.modulePath = fs::absolute(argv[1]);
415
416  if (!fs::exists(request.modulePath)) {
417    std::ostringstream msg;
418    msg << "Module not found: " << request.modulePath;
419    return absl::NotFoundError(msg.str());
420  }
421
422  if (request.modulePath.extension() != ".slang") {
423    std::ostringstream msg;
424    msg << "Expected a .slang file: " << request.modulePath;
425    return absl::InvalidArgumentError(msg.str());
426  }
427
428  request.moduleName = request.modulePath.stem().string();
429  request.searchPath = request.modulePath.has_parent_path()
430  ? request.modulePath.parent_path().string()
431  : fs::current_path().string();
432  request.outputPath = request.modulePath;
433  request.outputPath.replace_extension(".hlsl");
434
435  return request;
436}
437
438std::vector<CompilerOptionEntry> makeCommonOptions() {
439  std::vector<CompilerOptionEntry> options;
440  addCompilerOption(options, CompilerOptionName::DisableNonEssentialValidations);
441  addCompilerOption(options, CompilerOptionName::NoHLSLBinding);
442  addCompilerOption(options, CompilerOptionName::NoMangle);
443  addCompilerOption(options, CompilerOptionName::NoHLSLPackConstantBufferElements);
444  addCompilerOption(options, CompilerOptionName::PlainFunctionEntryPoints);
445  return options;
446}
447
448void configureTargetDesc(
449    IGlobalSession* globalSession,
450    std::vector<CompilerOptionEntry>& targetOptions,
451    TargetDesc& outDesc) {
452  outDesc = {};
453  outDesc.format = SLANG_HLSL;
454  outDesc.profile = globalSession->findProfile("lib_6_6");
455  outDesc.flags = SLANG_TARGET_FLAG_GENERATE_WHOLE_PROGRAM;
456  outDesc.compilerOptionEntries = targetOptions.data();
457  outDesc.compilerOptionEntryCount = static_cast<uint32_t>(targetOptions.size());
458}
459
460void configureSessionDesc(
461    const TargetDesc& targetDesc,
462    const ModuleRequest& request,
463    std::vector<CompilerOptionEntry>& sessionOptions,
464    std::array<const char*, 1>& searchPathStorage,
465    SessionDesc& outDesc) {
466  searchPathStorage[0] = request.searchPath.c_str();
467
468  outDesc = {};
469  outDesc.targets = &targetDesc;
470  outDesc.targetCount = 1;
471  outDesc.searchPaths = searchPathStorage.data();
472  outDesc.searchPathCount = static_cast<uint32_t>(searchPathStorage.size());
473  outDesc.compilerOptionEntries = sessionOptions.data();
474  outDesc.compilerOptionEntryCount = static_cast<uint32_t>(sessionOptions.size());
475}
476
477absl::StatusOr<ComPtr<IModule>> loadSlangModule(ISession* session, const std::string& moduleName) {
478  ComPtr<IModule> module;
479  ComPtr<IBlob> diagnostics;
480  module = session->loadModule(moduleName.c_str(), diagnostics.writeRef());
481
482  const std::string context = "loadModule: " + moduleName;
483  printDiagnostics(context.c_str(), diagnostics);
484
485  if (!module) {
486    std::ostringstream msg;
487    msg << "Failed to load module '" << moduleName << "'.";
488    return absl::InternalError(msg.str());
489  }
490
491  return module;
492}
493
494struct EntryPointsResult {
495  std::vector<FunctionInfo> functions;
496  std::unordered_set<std::string> publicFunctions;
497};
498
499absl::StatusOr<EntryPointsResult> collectEntryPoints(
500    IModule* module,
501    const std::string& moduleName,
502    const fs::path& sourcePath) {
503  std::vector<FunctionInfo> functions;
504  std::unordered_set<std::string> seenNames;
505
506  std::unordered_set<std::string> publicFunctions = findPublicFunctionNames(sourcePath);
507  if (publicFunctions.empty()) {
508    std::ostringstream msg;
509    msg << "No public functions found in " << sourcePath.string() << '.';
510    return absl::NotFoundError(msg.str());
511  }
512
513  DeclReflection* moduleReflection = module ? module->getModuleReflection() : nullptr;
514  if (!moduleReflection) {
515    std::ostringstream msg;
516    msg << "Failed to retrieve reflection data for module '"
517        << moduleName << "'.";
518    return absl::InternalError(msg.str());
519  }
520
521  collectFunctionInfos(moduleReflection, publicFunctions, functions, seenNames);
522
523  if (functions.empty()) {
524    std::ostringstream msg;
525    msg << "No public functions found in module '" << moduleName << "'.";
526    return absl::NotFoundError(msg.str());
527  }
528
529  return EntryPointsResult{functions, publicFunctions};
530}
531
532absl::StatusOr<ComPtr<ICompileRequest>> createCompileRequest(
533    ISession* session,
534    const ModuleRequest& request,
535    const TargetDesc& targetDesc,
536    const std::vector<FunctionInfo>& functions) {
537  ComPtr<ICompileRequest> compileRequest;
538  if (absl::Status status = checkSlangResult(
539          "ISession::createCompileRequest",
540          session->createCompileRequest(compileRequest.writeRef()));
541      !status.ok()) {
542    return status;
543  }
544
545  compileRequest->setCodeGenTarget(SLANG_HLSL);
546  compileRequest->setTargetProfile(0, targetDesc.profile);
547  compileRequest->setTargetFlags(0, SLANG_TARGET_FLAG_GENERATE_WHOLE_PROGRAM);
548  compileRequest->setMatrixLayoutMode(SLANG_MATRIX_LAYOUT_ROW_MAJOR);
549  compileRequest->setLineDirectiveMode(SLANG_LINE_DIRECTIVE_MODE_NONE);
550
551  compileRequest->addSearchPath(request.searchPath.c_str());
552
553  const int translationUnitIndex = compileRequest->addTranslationUnit(
554      SLANG_SOURCE_LANGUAGE_SLANG,
555      request.moduleName.c_str());
556  compileRequest->addTranslationUnitSourceFile(
557      translationUnitIndex,
558      request.modulePath.string().c_str());
559
560  for (const FunctionInfo& func : functions) {
561    const int entryPointIndex = compileRequest->addEntryPoint(
562        translationUnitIndex,
563        func.name.c_str(),
564        SLANG_STAGE_DISPATCH);
565    if (entryPointIndex < 0) {
566      std::ostringstream msg;
567      msg << "Failed to register entry point '" << func.name << "'.";
568      return absl::InternalError(msg.str());
569    }
570  }
571
572  return compileRequest;
573}
574
575absl::StatusOr<std::string> collectGeneratedHlsl(ICompileRequest* compileRequest, const std::string& moduleName) {
576  SlangResult compileResult = compileRequest->compile();
577  ComPtr<IBlob> diagnostics;
578  compileRequest->getDiagnosticOutputBlob(diagnostics.writeRef());
579  if (absl::Status status = checkSlangResult(
580          "ICompileRequest::compile", compileResult, diagnostics.get());
581      !status.ok()) {
582    return status;
583  }
584
585  ComPtr<IBlob> targetCodeBlob;
586  if (absl::Status status = checkSlangResult(
587          "ICompileRequest::getTargetCodeBlob",
588          compileRequest->getTargetCodeBlob(0, targetCodeBlob.writeRef()));
589      !status.ok()) {
590    return status;
591  }
592
593  if (!targetCodeBlob || targetCodeBlob->getBufferSize() == 0) {
594    std::ostringstream msg;
595    msg << "No HLSL was generated for module '" << moduleName << "'.";
596    return absl::InternalError(msg.str());
597  }
598
599  return std::string(
600      static_cast<const char*>(targetCodeBlob->getBufferPointer()),
601      static_cast<std::size_t>(targetCodeBlob->getBufferSize()));
602}
603
604std::string removeNvapiInclude(std::string hlslSource) {
605  const std::string guardToken = "#ifdef SLANG_HLSL_ENABLE_NVAPI";
606  const std::string endifToken = "#endif";
607
608  std::size_t searchPos = 0;
609  while (true) {
610    const std::size_t blockStart = hlslSource.find(guardToken, searchPos);
611    if (blockStart == std::string::npos) {
612      break;
613    }
614
615    std::size_t blockEnd = hlslSource.find(endifToken, blockStart);
616    if (blockEnd == std::string::npos) {
617      break;
618    }
619    blockEnd += endifToken.size();
620
621    while (blockEnd < hlslSource.size() &&
622        (hlslSource[blockEnd] == '\r' || hlslSource[blockEnd] == '\n')) {
623      ++blockEnd;
624    }
625
626    hlslSource.erase(blockStart, blockEnd - blockStart);
627    searchPos = blockStart;
628  }
629
630  return hlslSource;
631}
632
633
634std::string applyIncludeGuard(const std::string& hlslSource, const IncludeGuardInfo& includeGuard) {
635  if (!includeGuard.present) {
636    return hlslSource;
637  }
638
639  const std::string guardIfndefToken = "#ifndef " + includeGuard.macro;
640  const std::string guardDefineToken = "#define " + includeGuard.macro;
641  const bool alreadyGuarded =
642  hlslSource.find(guardIfndefToken) != std::string::npos &&
643  hlslSource.find(guardDefineToken) != std::string::npos;
644
645  if (alreadyGuarded) {
646    return hlslSource;
647  }
648
649  std::string body = hlslSource;
650  if (!body.empty() && body.back() != '\n') {
651    body += '\n';
652  }
653
654  std::ostringstream wrapped;
655  wrapped << includeGuard.ifndefLine << '\n';
656  wrapped << includeGuard.defineLine << '\n';
657  wrapped << '\n';
658  wrapped << body;
659  if (!body.empty() && body.back() != '\n') {
660    wrapped << '\n';
661  }
662  wrapped << includeGuard.endifLine;
663  if (!includeGuard.endifLine.empty() && includeGuard.endifLine.back() != '\n') {
664    wrapped << '\n';
665  }
666
667  return wrapped.str();
668}
669
670absl::Status run(int argc, char** argv) {
671  absl::StatusOr<ModuleRequest> requestOr = parseModuleRequest(argc, argv);
672  if (!requestOr.ok()) {
673    return requestOr.status();
674  }
675  ModuleRequest request = std::move(requestOr).value();
676
677  ComPtr<IGlobalSession> globalSession;
678  if (absl::Status status = checkSlangResult(
679          "createGlobalSession",
680          createGlobalSession(globalSession.writeRef()));
681      !status.ok()) {
682    return status;
683  }
684
685  auto commonOptions = makeCommonOptions();
686
687  std::vector<CompilerOptionEntry> targetOptions = commonOptions;
688  TargetDesc targetDesc;
689  configureTargetDesc(globalSession.get(), targetOptions, targetDesc);
690
691  std::vector<CompilerOptionEntry> sessionOptions = commonOptions;
692  SessionDesc sessionDesc;
693  std::array<const char*, 1> searchPaths{};
694  configureSessionDesc(targetDesc, request, sessionOptions, searchPaths, sessionDesc);
695
696  ComPtr<ISession> session;
697  if (absl::Status status = checkSlangResult(
698          "IGlobalSession::createSession",
699          globalSession->createSession(sessionDesc, session.writeRef()));
700      !status.ok()) {
701    return status;
702  }
703
704  absl::StatusOr<ComPtr<IModule>> libraryModuleOr =
705      loadSlangModule(session.get(), request.moduleName);
706  if (!libraryModuleOr.ok()) {
707    return libraryModuleOr.status();
708  }
709  ComPtr<IModule> libraryModule = std::move(libraryModuleOr).value();
710
711  absl::StatusOr<EntryPointsResult> entryPointsOr =
712      collectEntryPoints(libraryModule.get(), request.moduleName, request.modulePath);
713  if (!entryPointsOr.ok()) {
714    return entryPointsOr.status();
715  }
716  EntryPointsResult entryPoints = std::move(entryPointsOr).value();
717
718  absl::StatusOr<ComPtr<ICompileRequest>> compileRequestOr =
719      createCompileRequest(session.get(), request, targetDesc, entryPoints.functions);
720  if (!compileRequestOr.ok()) {
721    return compileRequestOr.status();
722  }
723  ComPtr<ICompileRequest> compileRequest = std::move(compileRequestOr).value();
724
725  absl::StatusOr<std::string> hlslSourceOr =
726      collectGeneratedHlsl(compileRequest.get(), request.moduleName);
727  if (!hlslSourceOr.ok()) {
728    return hlslSourceOr.status();
729  }
730  std::string hlslSource = std::move(hlslSourceOr).value();
731  std::string filteredHlsl = removeNvapiInclude(hlslSource);
732
733  fs::path rawOutputPath = request.outputPath;
734  rawOutputPath.replace_extension(".raw.hlsl");
735  if (absl::Status status = writeTextFile(rawOutputPath, hlslSource);
736      !status.ok()) {
737    return status;
738  }
739
740  IncludeGuardInfo includeGuard = detectIncludeGuard(request.modulePath);
741  std::string finalHlsl = applyIncludeGuard(filteredHlsl, includeGuard);
742  if (absl::Status status = writeTextFile(request.outputPath, finalHlsl);
743      !status.ok()) {
744    return status;
745  }
746
747  std::cerr << "Generated HLSL written to " << request.outputPath << std::endl;
748  return absl::OkStatus();
749}
750
751int main(int argc, char** argv) {
752  absl::Status status = run(argc, argv);
753  if (!status.ok()) {
754    std::cerr << status.message() << std::endl;
755    return 1;
756  }
757
758  return 0;
759}