Add support for heterogeneous compute (#2836)

* Extract device code from assembly files.

This commit by itself does not actually add any viewing panes for the
extracted device code, nor does it provide any means to provide
device-specific postprocessing steps.

* Add a rudimentary device view.

This doesn't handle multiple device outputs correctly at the moment, and
it provides no opportunity to view device code differently from
non-device code for now.

* Lint issues

* Some more WIP getting it working.

* Add device extraction code

* More fixes, and keep state

Co-authored-by: Joshua Cranmer <joshua.cranmer@intel.com>
Co-authored-by: Patrick Quist <partouf@gmail.com>

1 files changed, 76 insertions, 3 deletions

diff --git a/lib/compilers/clang.js b/lib/compilers/clang.js
index 5e3252965..b849ef54f 100644
--- a/lib/compilers/clang.js
+++ b/lib/compilers/clang.js
@@ -28,8 +28,17 @@ import { AmdgpuAsmParser } from '../asm-parser-amdgpu';
 import { SassAsmParser } from '../asm-parser-sass';
 import { BaseCompiler } from '../base-compiler';
 
+const offloadRegexp = /^#\s+__CLANG_OFFLOAD_BUNDLE__(__START__|__END__)\s+(.*)$/gm;
+
 export class ClangCompiler extends BaseCompiler {
-    static get key() { return 'clang'; }
+    static get key() {
+        return 'clang';
+    }
+
+    constructor(info, env) {
+        super(info, env);
+        this.compiler.supportsDeviceAsmView = true;
+    }
 
     runCompiler(compiler, options, inputFilename, execOptions) {
         if (!execOptions) {
@@ -40,10 +49,70 @@ export class ClangCompiler extends BaseCompiler {
 
         return super.runCompiler(compiler, options, inputFilename, execOptions);
     }
+
+    splitDeviceCode(assembly) {
+        // Check to see if there is any offload code in the assembly file.
+        if (!offloadRegexp.test(assembly))
+            return null;
+
+        offloadRegexp.lastIndex = 0;
+        const matches = assembly.matchAll(offloadRegexp);
+        let prevStart = 0;
+        const devices = {};
+        for (const match of matches) {
+            const [full, startOrEnd, triple] = match;
+            if (startOrEnd === '__START__') {
+                prevStart = match.index + full.length + 1;
+            } else {
+                devices[triple] = assembly.substr(prevStart, match.index - prevStart);
+            }
+        }
+        return devices;
+    }
+
+    extractDeviceCode(result, filters) {
+        const split = this.splitDeviceCode(result.asm);
+        if (!split)
+            return result;
+
+        const devices = result.devices = {};
+        for (const key of Object.keys(split)) {
+            if (key.indexOf('host-') === 0)
+                result.asm = split[key];
+            else
+                devices[key] = this.processDeviceAssembly(key, split[key], filters);
+        }
+        return result;
+        // //result.asm = ...
+        //
+        // const extractor = this.compiler.deviceExtractor;
+        // const extractor_result = await this.exec(extractor,
+        //     [outputFilename, result.dirPath + '/devices'], {});
+        // if (extractor_result.code !== 0) {
+        //     result.devices = extractor_result.stderr;
+        //     return result;
+        // }
+        // const extractor_info = JSON.parse(extractor_result.stdout);
+        // result.asm = await fs.readFile(extractor_info.host, 'utf-8');
+        // const devices = result.devices = {};
+        // for (const device in extractor_info.devices) {
+        //     const file = extractor_info.devices[device];
+        //     devices[device] = await this.processDeviceAssembly(device, file, filters);
+        // }
+        // return result;
+    }
+
+    processDeviceAssembly(deviceName, deviceAsm, filters) {
+        return this.llvmIr.isLlvmIr(deviceAsm) ?
+            this.llvmIr.process(deviceAsm, filters) : this.asm.process(deviceAsm, filters);
+    }
 }
 
 export class ClangCudaCompiler extends ClangCompiler {
-    static get key() { return 'clang-cuda'; }
+    static get key() {
+        return 'clang-cuda';
+    }
+
     constructor(info, env) {
         super(info, env);
 
@@ -69,8 +138,12 @@ export class ClangCudaCompiler extends ClangCompiler {
         return result;
     }
 }
+
 export class ClangHipCompiler extends ClangCompiler {
-    static get key() { return 'clang-hip'; }
+    static get key() {
+        return 'clang-hip';
+    }
+
     constructor(info, env) {
         super(info, env);