Home
last modified time | relevance | path

Searched refs:gpu_module (Results 1 – 12 of 12) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tools/kernel_gen/transforms/
H A Dgpu_kernel_to_blob_pass.cc62 gpu::GPUModuleOp gpu_module = getOperation(); in runOnOperation() local
63 auto blob_or = GetGpuBinaryBlob(gpu_module); in runOnOperation()
67 gpu_module->setAttr(blob_annotation_, in runOnOperation()
72 gpu_module.emitError(blob_or.status().error_message()); in runOnOperation()
77 gpu::GPUModuleOp gpu_module) { in GetGpuBinaryBlob() argument
85 auto llvmModule = translateModuleToLLVMIR(gpu_module, llvmContext); in GetGpuBinaryBlob()
90 llvmModule->setModuleIdentifier(gpu_module.getName()); in GetGpuBinaryBlob()
170 << gpu_module.getName() << "' on architecture sm_" << arch in GetGpuBinaryBlob()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/tests/lhlo_to_jitrt/
H A Dfusion_to_gpu_binary.mlir20 // CHECK: gpu.module @[[gpu_module:.*]] attributes {
38 // CHECK: gpu.launch_func @[[gpu_module]]::@[[kernel]]
61 // CHECK: gpu.module @[[gpu_module:.*]] attributes {
87 // CHECK: gpu.launch_func @[[gpu_module]]::@[[kernel]]
H A Dlaunch_to_jitrt.mlir23 gpu.module @gpu_module attributes {binary = "kernel binary"} {
52 gpu.launch_func @gpu_module::@main
H A Dpad_to_static_to_gpu_binary.mlir19 // CHECK: gpu.module @gpu_module
H A Dslice_to_dynamic_to_gpu_binary.mlir19 // CHECK: gpu.module @gpu_module
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/
H A Dkernel_ops_pattern.cc341 auto gpu_module = rewriter.create<mlir::gpu::GPUModuleOp>(loc, "gpu_module"); in Rewrite() local
342 symbol_table.insert(gpu_module); in Rewrite()
343 gpu_module->setAttr(tfrt::gpu::GetGpuBinaryAttrName(), in Rewrite()
358 mlir::SymbolRefAttr::get(gpu_module)); in Rewrite()
363 gpu_module->setAttr(tfrt::gpu::GetGpuConstantsAttrName(), in Rewrite()
410 rewriter.setInsertionPointToStart(gpu_module.getBody()); in Rewrite()
H A Dlmhlo_to_jitrt.cc373 auto gpu_module = kernel->getParentOfType<mlir::gpu::GPUModuleOp>(); in matchAndRewrite() local
374 auto gpu_binary = gpu_module->getAttrOfType<mlir::StringAttr>("binary"); in matchAndRewrite()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/
H A Dpropagate_static_shapes.mlir7 gpu.module @gpu_module {
30 gpu.launch_func @gpu_module::@kernel
/aosp_15_r20/external/pytorch/test/
H A Dtest_modules.py570 gpu_module = module_cls(*args, **kwargs).to(dtype).to(device)
571 gpu_module.train(training)
577 gpu_module(*gpu_forward_args, **gpu_forward_kwargs)
579 for cpu_p, gpu_p in zip(cpu_module.parameters(), gpu_module.parameters()):
584 gpu_outputs = gpu_module(*gpu_forward_args, **gpu_forward_kwargs)
600 for cpu_p, gpu_p in zip(cpu_module.parameters(), gpu_module.parameters()):
/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcommon_nn.py3519 gpu_module = self.constructor(*self.constructor_args).float().cuda()
3521 gpu_param = test_case._get_parameters(gpu_module)
3528 test_case._zero_grad_parameters(gpu_module)
3530 gpu_output = test_case._forward(gpu_module, gpu_input_tuple)
3541 … gpu_gradInput = test_case._backward(gpu_module, gpu_input_tuple, gpu_output, gpu_gradOutput)
3549 gpu_output = gpu_module(*gpu_input_tuple)
3565 gpu_input_tuple + tuple(gpu_module.parameters()),
3589 gpu_input_tuple + (gpu_gradOutput,) + tuple(gpu_module.parameters()),
3595 self.test_noncontig(test_case, gpu_module, gpu_input_tuple)
3887 gpu_module = self.constructor(*self.constructor_args)
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tools/kernel_gen/
H A Dkernel_creator.cc318 for (::mlir::gpu::GPUModuleOp gpu_module : gpu_modules) { in LowerKernelBodiesToLowLevelIr()
319 gpu_module.walk([&](mlir::gpu::GPUFuncOp gpu_kernel) { in LowerKernelBodiesToLowLevelIr()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dir_emitter_unnested.cc1600 auto gpu_module = kernel_func->getParentOfType<mlir::gpu::GPUModuleOp>(); in EmitLaunchFunc() local
1602 gpu_module, module_->getContext(), gpu_module.getName()); in EmitLaunchFunc()