From 28231a718fd6ee02261cf48322625e2218a0ad66 Mon Sep 17 00:00:00 2001
From: Leon Chen <leon@md.ai>
Date: Mon, 27 Nov 2017 13:36:11 -0500
Subject: [PATCH] clean up glTextureFragmentsAsColStack after use to prevent
 downstream effects on running shader programs on fragments

---
 src/Tensor.js                               | 12 ++++++++++++
 src/layers/convolutional/Conv2D.js          |  3 +++
 src/layers/convolutional/Conv2DTranspose.js |  4 +++-
 src/layers/convolutional/Conv3D.js          |  3 +++
 src/layers/convolutional/SeparableConv2D.js |  3 +++
 src/layers/core/Flatten.js                  |  2 +-
 src/layers/pooling/_Pooling1D.js            |  1 +
 src/layers/pooling/_Pooling2D.js            |  1 +
 src/layers/pooling/_Pooling3D.js            |  1 +
 9 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/Tensor.js b/src/Tensor.js
index 4648439..1bc454f 100644
--- a/src/Tensor.js
+++ b/src/Tensor.js
@@ -227,6 +227,18 @@ export default class Tensor {
     gl.deleteFramebuffer(fbo)
   }
 
+  /**
+   * Removes glTextureFragmentsAsColStack
+   */
+  removeGLTextureFragmentsAsColStack() {
+    if (this.glTextureFragmentsAsColStack) {
+      const gl = webgl2.context
+      gl.deleteTexture(this.glTextureFragmentsAsColStack)
+      delete this.glTextureFragmentsAsColStack
+      delete this.glTextureFragmentsAsColStackShape
+    }
+  }
+
   /**
    * Deletes WebGLTexture
    */
diff --git a/src/layers/convolutional/Conv2D.js b/src/layers/convolutional/Conv2D.js
index d638869..553c97e 100644
--- a/src/layers/convolutional/Conv2D.js
+++ b/src/layers/convolutional/Conv2D.js
@@ -445,6 +445,9 @@ export default class Conv2D extends Layer {
         ],
         supportsTextureFragments: true
       })
+      if (hasFragments) {
+        x.removeGLTextureFragmentsAsColStack()
+      }
     } else {
       // run matrix multiply on result of im2col
       const matMulInputs = [{ input: this.imColsMat, name: 'A' }, { input: this.weights['kernel'], name: 'B' }]
diff --git a/src/layers/convolutional/Conv2DTranspose.js b/src/layers/convolutional/Conv2DTranspose.js
index 473942f..29e7b90 100644
--- a/src/layers/convolutional/Conv2DTranspose.js
+++ b/src/layers/convolutional/Conv2DTranspose.js
@@ -471,7 +471,6 @@ export default class Conv2DTranspose extends Layer {
       )
       this.convTransposeProgram = webgl2.compileProgram(convTransposeProgramSource)
     }
-
     webgl2.runProgram({
       program: this.convTransposeProgram,
       output: this.activation === 'linear' ? this.output : this.outputPreactiv,
@@ -482,6 +481,9 @@ export default class Conv2DTranspose extends Layer {
       ],
       supportsTextureFragments: true
     })
+    if (hasFragments) {
+      this.matMulResult.removeGLTextureFragmentsAsColStack()
+    }
 
     // Activation
     if (this.activation !== 'linear') {
diff --git a/src/layers/convolutional/Conv3D.js b/src/layers/convolutional/Conv3D.js
index 2fe18cb..aaee0ac 100644
--- a/src/layers/convolutional/Conv3D.js
+++ b/src/layers/convolutional/Conv3D.js
@@ -503,6 +503,9 @@ export default class Conv3D extends Layer {
         ],
         supportsTextureFragments: true
       })
+      if (hasFragments) {
+        x.removeGLTextureFragmentsAsColStack()
+      }
     } else {
       // run matrix multiply on result of vol2col
       const matMulInputs = [{ input: this.volColsMat, name: 'A' }, { input: this.weights['kernel'], name: 'B' }]
diff --git a/src/layers/convolutional/SeparableConv2D.js b/src/layers/convolutional/SeparableConv2D.js
index 7a210d0..8d17fe1 100644
--- a/src/layers/convolutional/SeparableConv2D.js
+++ b/src/layers/convolutional/SeparableConv2D.js
@@ -168,6 +168,9 @@ class _DepthwiseConv2D extends Conv2D {
       uniforms: [{ value: this.output.glTextureShape[1], type: 'int', name: 'inputCols' }],
       supportsTextureFragments: true
     })
+    if (this.output.glTextureFragments) {
+      this.output.removeGLTextureFragmentsAsColStack()
+    }
   }
 }
 
diff --git a/src/layers/core/Flatten.js b/src/layers/core/Flatten.js
index 6385a05..fb63f4f 100644
--- a/src/layers/core/Flatten.js
+++ b/src/layers/core/Flatten.js
@@ -76,7 +76,6 @@ export default class Flatten extends Layer {
 
     if (x.glTextureFragments) {
       x.convert2DRowFragmentedGLTextureToColStack()
-      console.log(x)
       webgl2.runProgram({
         program: this.flattenFragmentsProgram,
         output: this.output,
@@ -88,6 +87,7 @@ export default class Flatten extends Layer {
         ],
         supportsTextureFragments: true
       })
+      x.removeGLTextureFragmentsAsColStack()
     } else {
       webgl2.runProgram({
         program: this.flattenProgram,
diff --git a/src/layers/pooling/_Pooling1D.js b/src/layers/pooling/_Pooling1D.js
index bf7d55b..dcc39d5 100644
--- a/src/layers/pooling/_Pooling1D.js
+++ b/src/layers/pooling/_Pooling1D.js
@@ -176,6 +176,7 @@ export default class _Pooling1D extends Layer {
         uniforms: programUniforms,
         supportsTextureFragments: true
       })
+      x.removeGLTextureFragmentsAsColStack()
     } else {
       webgl2.runProgram({
         program: this.poolingProgram,
diff --git a/src/layers/pooling/_Pooling2D.js b/src/layers/pooling/_Pooling2D.js
index 888826e..8469816 100644
--- a/src/layers/pooling/_Pooling2D.js
+++ b/src/layers/pooling/_Pooling2D.js
@@ -318,6 +318,7 @@ export default class _Pooling2D extends Layer {
         uniforms: programUniforms,
         supportsTextureFragments: true
       })
+      input.removeGLTextureFragmentsAsColStack()
     } else {
       webgl2.runProgram({
         program: this.poolingProgram,
diff --git a/src/layers/pooling/_Pooling3D.js b/src/layers/pooling/_Pooling3D.js
index ca65293..bd8acd2 100644
--- a/src/layers/pooling/_Pooling3D.js
+++ b/src/layers/pooling/_Pooling3D.js
@@ -383,6 +383,7 @@ export default class _Pooling3D extends Layer {
         uniforms: programUniforms,
         supportsTextureFragments: true
       })
+      input.removeGLTextureFragmentsAsColStack()
     } else {
       webgl2.runProgram({
         program: this.poolingProgram,