| 33 | namespace { |
| 34 | |
| 35 | class Add : public NodeShader { |
| 36 | public: |
| 37 | Status GenerateCode(const GenerationContext& ctx, |
| 38 | GeneratedCode* generated_code) const final { |
| 39 | auto attr = absl::any_cast<AddAttributes>(ctx.node->operation.attributes); |
| 40 | auto adds = absl::get_if<Tensor<Linear, DataType::FLOAT32>>(&attr.param); |
| 41 | auto scalar = absl::get_if<float>(&attr.param); |
| 42 | auto inputs = ctx.graph->FindInputs(ctx.node->id); |
| 43 | |
| 44 | if (!adds && !scalar) { |
| 45 | // check if it is a broadcast |
| 46 | if (inputs.size() == 2 && |
| 47 | inputs[0]->tensor.shape != inputs[1]->tensor.shape && |
| 48 | inputs[1]->tensor.shape.h == 1 && inputs[1]->tensor.shape.w == 1 && |
| 49 | inputs[0]->tensor.shape.c == inputs[1]->tensor.shape.c) { |
| 50 | *generated_code = { |
| 51 | /*parameters=*/{}, |
| 52 | /*objects=*/{}, |
| 53 | /*shared_variables=*/{}, |
| 54 | /*workload=*/uint3(), |
| 55 | /*workgroup=*/uint3(), |
| 56 | /*source_code=*/ |
| 57 | "value_0 = $input_data_1[gid.z]$ + $input_data_0[gid.x, gid.y, " |
| 58 | "gid.z]$;", |
| 59 | /*input=*/IOStructure::ONLY_DEFINITIONS, |
| 60 | /*output=*/IOStructure::AUTO, |
| 61 | }; |
| 62 | return OkStatus(); |
| 63 | } |
| 64 | |
| 65 | std::string code = "value_0 = value_0"; |
| 66 | for (int index = 1; index < inputs.size(); ++index) { |
| 67 | if (inputs[index]->tensor.shape != inputs[0]->tensor.shape) { |
| 68 | return InvalidArgumentError("Shapes are not equal"); |
| 69 | } |
| 70 | absl::StrAppend(&code, " + value_", index); |
| 71 | } |
| 72 | absl::StrAppend(&code, ";"); |
| 73 | *generated_code = { |
| 74 | /*parameters=*/{}, |
| 75 | /*objects=*/{}, |
| 76 | /*shared_variables=*/{}, |
| 77 | /*workload=*/uint3(), |
| 78 | /*workgroup=*/uint3(), |
| 79 | /*source_code=*/std::move(code), |
| 80 | /*input=*/IOStructure::AUTO, |
| 81 | /*output=*/IOStructure::AUTO, |
| 82 | }; |
| 83 | return OkStatus(); |
| 84 | } |
| 85 | |
| 86 | if (scalar) { |
| 87 | *generated_code = { |
| 88 | /*parameters=*/{{"scalar", *scalar}}, |
| 89 | /*objects=*/{}, |
| 90 | /*shared_variables=*/{}, |
| 91 | /*workload=*/uint3(), |
| 92 | /*workgroup=*/uint3(), |