| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_dynamic_padding_mapper.cc | 58 tf_device::LaunchFuncOp launch_func, tf_device::ReplicateOp replicate) { in GetRemappedReplicatedInputIndices() argument
62 for (auto operand_and_idx : llvm::enumerate(launch_func.getOperands())) in GetRemappedReplicatedInputIndices()
74 tf_device::LaunchFuncOp launch_func, int num_replicated_args, in GetRemappedPaddings() argument
87 Attribute padding_map_attr = launch_func.getAttr(kPaddingMapAttr); in GetRemappedPaddings()
92 return launch_func.emitOpError() in GetRemappedPaddings()
100 return launch_func.emitOpError( in GetRemappedPaddings()
106 return launch_func.emitOpError(llvm::formatv( in GetRemappedPaddings()
113 return launch_func.emitOpError() in GetRemappedPaddings()
118 return launch_func.emitOpError() in GetRemappedPaddings()
127 return launch_func.emitOpError(llvm::formatv( in GetRemappedPaddings()
[all …]
|
| D | tpu_rewrite_pass.cc | 254 Operation* BuildCompileOp(tf_device::LaunchFuncOp launch_func, int num_replicas, in BuildCompileOp() argument
259 OperationState compile_op_state(launch_func.getLoc(), "tf._TPUCompileMlir"); in BuildCompileOp()
264 launch_func, num_replicas, num_cores_per_replica, &metadata))) in BuildCompileOp()
281 compile_op_operands.reserve(launch_func.getNumOperands()); in BuildCompileOp()
283 for (auto operand_and_idx : llvm::enumerate(launch_func.getOperands())) { in BuildCompileOp()
290 launch_func.getLoc(), in BuildCompileOp()
301 launch_func.getAttrOfType<FlatSymbolRefAttr>("func"); in BuildCompileOp()
303 launch_func.emitOpError("does not have `func` attribute"); in BuildCompileOp()
306 FuncOp func = launch_func.getParentOfType<ModuleOp>().lookupSymbol<FuncOp>( in BuildCompileOp()
331 tf_device::LaunchFuncOp launch_func, in BuildExecuteOp() argument
[all …]
|
| D | annotate_parameter_replication.cc | 59 m.walk([&](tf_device::LaunchFuncOp launch_func) { in runOnModule() argument
60 auto replicate = launch_func.getParentOfType<tf_device::ReplicateOp>(); in runOnModule()
71 auto func = llvm::cast<FuncOp>(m.lookupSymbol(launch_func.func())); in runOnModule()
72 for (auto entry : llvm::enumerate(launch_func.getOperands())) { in runOnModule()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/ |
| D | tpu_dynamic_padding_mapper.mlir | 13 …"tf_device.launch_func"(%ri_0, %ri_1) {device = "", func = @func0, padding_map = ["\10\02\18\01"]}…
40 …"tf_device.launch_func"(%ri_0, %ri_1, %ri_2) {device = "", func = @func1, padding_map = ["\10\02\1…
72 …"tf_device.launch_func"(%ri_0, %ri_1, %ri_2, %ri_3, %ri_4) {device = "", func = @func2, padding_ma…
93 …"tf_device.launch_func"(%ri_1, %arg0, %ri_0) {device = "", func = @func3, padding_map = ["\10\02\1…
114 …"tf_device.launch_func"(%arg0, %arg0, %arg0) {device = "", func = @func4, padding_map = ["\10\02\1…
128 …"tf_device.launch_func"(%ri_1, %arg0, %ri_0) {device = "", func = @func5} : (tensor<i1>, tensor<i1…
144 …"tf_device.launch_func"(%ri_1, %arg0, %ri_0) {device = "", func = @func6, padding_map = []} : (ten…
165 …"tf_device.launch_func"(%ri_1) {device = "", func = @func7, padding_map = ["\10\02\18\01"]} : (ten…
182 // expected-error@+1 {{'tf_device.launch_func' op requires 'padding_map' array attribute}}
183 "tf_device.launch_func"() {device = "", func = @_func, padding_map = 0 : i32} : () -> ()
[all …]
|
| D | tpu_rewrite.mlir | 8 …"tf_device.launch_func"() {_tpu_replicate = "cluster0", device = "", func = @empty_func, num_cores…
23 …"tf_device.launch_func"() {_tpu_replicate = "cluster0", device = "", func = @empty_func, num_cores…
33 // Tests `tf_device.launch_func` with missing `num_cores_per_replicas`
39 …"tf_device.launch_func"() {_tpu_replicate = "cluster0", device = "", func = @empty_func, step_mark…
49 // Tests `tf_device.launch_func` with bad `num_cores_per_replicas` attribute.
54 …"tf_device.launch_func"() {_tpu_replicate = "cluster0", device = "", func = @empty_func, num_cores…
64 // Tests `tf_device.launch_func` with missing `step_marker_location` attribute.
69 …"tf_device.launch_func"() {_tpu_replicate = "cluster0", device = "", func = @empty_func, num_cores…
79 // Tests `tf_device.launch_func` with bad `step_marker_location` attribute.
84 …"tf_device.launch_func"() {_tpu_replicate = "cluster0", device = "", func = @empty_func, num_cores…
[all …]
|
| D | cluster_outlining.mlir | 14 …// CHECK: %[[C_OUTPUT:[0-9]*]] = "tf_device.launch_func"(%[[A_OUTPUT]]) {device = "tpu0", func = @…
48 …// CHECK: %[[C_OUTPUT:[0-9]*]] = "tf_device.launch_func"(%[[A_OUTPUT]]) {device = "tpu0", func = @…
57 …// CHECK: %[[E_OUTPUT:[0-9]*]] = "tf_device.launch_func"(%[[C_OUTPUT]], %[[D_OUTPUT]]) {device = "…
94 …// CHECK: %[[A_OUTPUT:[a-z0-9]*]], %{{.*}} = {{.*}} "tf_device.launch_func"() {device = "tpu0", fu…
113 // Tests launch attributes are copied over to launch_func.
|
| D | annotate-parameter-replication.mlir | 13 …%4 = "tf_device.launch_func"(%ri_0, %3, %2) {func = @tpu0_func, device = ""} : (tensor<?xi32>, ten…
49 …%2 = "tf_device.launch_func"(%0, %1, %ri_2) {func = @tpu0_func, device = ""} : (tensor<?xi32>, ten…
75 …%2 = "tf_device.launch_func"(%0, %1) {func = @tpu0_func, device = ""} : (tensor<?xi32>, tensor<?xi…
|
| /external/tensorflow/tensorflow/compiler/xla/service/mlir_gpu/tests/ |
| D | mlir_gpu_lhlo_gen_test.cc | 144 ;CHECK: "gpu.launch_func"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[ARG0]], %[[ARG1]… in TEST_F()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/ |
| D | tf_device_ops.td | 106 def TfDevice_LaunchFuncOp : TfDevice_Op<"launch_func", []> {
108 The `tf_device.launch_func` launches a function on target device.
|