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 D | gpu_kernel_to_blob_pass.cc | 62 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 D | fusion_to_gpu_binary.mlir | 20 // 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 D | launch_to_jitrt.mlir | 23 gpu.module @gpu_module attributes {binary = "kernel binary"} { 52 gpu.launch_func @gpu_module::@main
|
H A D | pad_to_static_to_gpu_binary.mlir | 19 // CHECK: gpu.module @gpu_module
|
H A D | slice_to_dynamic_to_gpu_binary.mlir | 19 // CHECK: gpu.module @gpu_module
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/ |
H A D | kernel_ops_pattern.cc | 341 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 D | lmhlo_to_jitrt.cc | 373 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 D | propagate_static_shapes.mlir | 7 gpu.module @gpu_module { 30 gpu.launch_func @gpu_module::@kernel
|
/aosp_15_r20/external/pytorch/test/ |
H A D | test_modules.py | 570 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 D | common_nn.py | 3519 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 D | kernel_creator.cc | 318 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 D | ir_emitter_unnested.cc | 1600 auto gpu_module = kernel_func->getParentOfType<mlir::gpu::GPUModuleOp>(); in EmitLaunchFunc() local 1602 gpu_module, module_->getContext(), gpu_module.getName()); in EmitLaunchFunc()
|