gx

chenyc

2025-02-12 ea42ff3ebee1eeb3fb29423aa848a249441db81c

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
/**
 * @license
 * Copyright 2017 Google Inc. All Rights Reserved.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * =============================================================================
 */
 
import * as broadcast_util from '../../ops/broadcast_util';
 
import {GPGPUProgram} from './gpgpu_math';
 
const CHECK_NAN_SNIPPET = `
  if (isnan(a)) return a;
  if (isnan(b)) return b;
`;
 
export const ADD = 'return a + b;';
export const SUB = 'return a - b;';
export const MUL = 'return a * b;';
 
// Without the equality check div produces 0.9999 for a = b, which when
// floored can cause errors.
export const DIV = `
if (a == b) {
  return 1.0;
};
return a / b;`;
 
// We use native integer division to deal with floating point imprecision. Since
// we implement floor division and glsl implements truncated division, we
// correct for this by subtracting 1 from result when the result is negative and
// there is a remainder.
export const INT_DIV = `
  float s = sign(a) * sign(b);
  int ia = round(a);
  int ib = round(b);
  if (ib != 0) {
    // Windows (D3D) wants guaranteed non-zero int division at compile-time.
    return float(idiv(ia, ib, s));
  } else {
    return NAN;
  }
`;
 
export const POW = `
if(a < 0.0 && floor(b) < b){
  return NAN;
}
if (b == 0.0) {
  return 1.0;
}
return (round(mod(b, 2.0)) != 1) ?
    pow(abs(a), b) : sign(a) * pow(abs(a), b);
`;
export const SQUARED_DIFFERENCE = 'return (a - b) * (a - b);';
 
export const EQUAL = `return float(a == b);`;
 
export const NOT_EQUAL = `return float(a != b);`;
 
export const LESS = `return float(a < b);`;
 
export const LESS_EQUAL = `return float(a <= b);`;
 
export const GREATER = `return float(a > b);`;
 
export const GREATER_EQUAL = `return float(a >= b);`;
 
export const LOGICAL_AND = `return float(a >= 1.0 && b >= 1.0);`;
 
export const LOGICAL_OR = `return float(a >= 1.0 || b >= 1.0);`;
 
export const MAX = CHECK_NAN_SNIPPET + `
  return max(a, b);
`;
export const MIN = CHECK_NAN_SNIPPET + `
  return min(a, b);
`;
export const MOD = `if (b == 0.0) return NAN;
  return mod(a, b);`;
 
export const ATAN2 = CHECK_NAN_SNIPPET + `
  return atan(a, b);
`;
 
export const ELU_DER = `return (b >= 1.0) ? a : a * (b + 1.0);`;
 
export const PRELU = `return (a < 0.) ? b * a : a;`;
 
export class BinaryOpProgram implements GPGPUProgram {
  variableNames = ['A', 'B'];
  outputShape: number[];
  userCode: string;
 
  constructor(op: string, aShape: number[], bShape: number[]) {
    this.outputShape =
        broadcast_util.assertAndGetBroadcastShape(aShape, bShape);
    this.userCode = `
      float binaryOperation(float a, float b) {
        ${op}
      }
 
      void main() {
        float a = getAAtOutCoords();
        float b = getBAtOutCoords();
        setOutput(binaryOperation(a, b));
      }
    `;
  }
}