def generate(cls, graph: Graph, **kwargs):
graph, _ = WebGPUOptimizeRule().optimize(graph)
if flags.DEBUG:
traverse.dump(graph)
memory_layout = allocate(graph)
console.debug(f"[WebGPUDescriptorGenerator] memory_layout total size: {memory_layout.total_size * 4}[B]")
console.debug(f"[WebGPUDescriptorGenerator] memory_layout static size: {memory_layout.static_size * 4}[B]")
console.debug(f"[WebGPUDescriptorGenerator] memory_layout dynamic size: {memory_layout.dynamic_size * 4}[B]")
constant_encoder = ConstantEncoder.get_encoder(kwargs.get("constant_encoder_name", None))
constants_bytes = constant_encoder.encode(memory_layout)
console.debug(f"[WebGPUDescriptorGenerator] constants encoded size: {len(constants_bytes)}[B]")
kernels = cls.generate_kernels(graph, memory_layout)
descriptor = GraphDescriptor(
kernels=kernels,
memory_layout=memory_layout,
inputs=graph.inputs,
outputs=graph.outputs,
constants_encoding=constant_encoder.name,
licenses=graph.licenses
)
if flags.optimize.VALIDATE_GENERATED_SOURCE:
validate_kernel_source(descriptor)
return GraphExecutionData(graph, descriptor, constants_bytes)
def validate_kernel_source(descriptor: GraphDescriptor):
# FIXME: WebGPU supports multi shader languages, but this test supposes the language as METAL.
source = descriptor.concat_kernel_sources()
if os.name != 'posix':
# os.name in mac is 'posix', and xcrun command is only in mac
console.warning(
"[WebGPUDescriptorGenerator] 'xcrun' command is not found. validation of generated source code in webgpu backend is "
"skipped.")
return
with tmp.TemporaryDirectory() as tmpdir:
source_path = path.join(tmpdir, "kernel.metal")
lib_path = path.join(tmpdir, "kernel.air")
with open(source_path, "w+") as f:
f.write(source)
try:
result = subprocess.run([
"xcrun", "-sdk", "macosx", "metal", source_path, "-o", lib_path
],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
if result.returncode == 0:
if result.stderr == b"":
console.debug(
"[WebGPUDescriptorGenerator] Generated kernel source is valid."
)
else:
console.warning(
"[WebGPUDescriptorGenerator] In validating kernel source, warnings are generated."
)
console.stderr(result.stderr.decode("utf-8"))
else:
console.error(
"[WebGPUDescriptorGenerator] Generated kernel source is invalid."
)
console.stderr(result.stderr.decode("utf-8"))
exit(result.returncode)
except FileNotFoundError:
console.warning(
"[WebGPUDescriptorGenerator] 'xcrun' command is not found. validation of generated source code in webgpu backend is "
"skipped.")
return
def validate_kernel_source(descriptor: GraphDescriptor):
# FIXME: WebGPU supports multi shader languages, but this test supposes the language as METAL.
source = descriptor.concat_kernel_sources()
with tmp.TemporaryDirectory() as tmpdir:
source_path = path.join(tmpdir, "kernel.metal")
lib_path = path.join(tmpdir, "kernel.air")
with open(source_path, "w+") as f:
f.write(source)
result = subprocess.run(
["xcrun", "-sdk", "macosx", "metal", source_path, "-o", lib_path])
if result.returncode != 0:
print("Generated kernel source is invalid.")
exit(result.returncode)
def generate(graph: Graph,
constant_encoder_name: str = None) -> GraphExecutionData:
graph, _ = WebGPUOptimizeRule().optimize(graph)
if flags.DEBUG:
traverse.dump(graph)
variables_layout, constants_layout, constants_data = Allocator.allocate(
graph)
constant_encoder = ConstantEncoder.get_encoder(constant_encoder_name)
constants_bytes = constant_encoder.encode(constants_layout, constants_data)
if flags.DEBUG:
print(
f"[GraphDescriptorGeneratorWebGPU] allocated constant-buffer size: {constants_layout.size * 4} [Byte]"
)
print(
f"[GraphDescriptorGeneratorWebGPU] encoded constant-buffer size: {len(constants_bytes)} [Byte]"
)
print(
f"[GraphDescriptorGeneratorWebGPU] allocated variable-buffer size: {variables_layout.size * 4} [Byte]"
)
kernels = generate_kernels(graph, constants_layout, variables_layout)
descriptor = GraphDescriptor(kernels=kernels,
constants_layout=constants_layout,
variables_layout=variables_layout,
inputs=graph.inputs,
outputs=graph.outputs,
constants_encoding=constant_encoder.name,
licenses=graph.licenses)
if flags.optimize.VALIDATE_GENERATED_SOURCE:
if flags.DEBUG:
print(
"[GraphDescriptorGeneratorWebGPU] validate generated kernel source"
)
validate_kernel_source(descriptor)
return GraphExecutionData(descriptor, constants_bytes)