Add winograd into annotation pass

apache · Jul 19, 2022 · be3076d · be3076d
1 parent 71f8f2d
commit be3076d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/relay/transforms/annotate_texture_storage.cc b/src/relay/transforms/annotate_texture_storage.cc
@@ -324,6 +324,11 @@ class StorageInfo : private transform::DeviceAwareExprVisitor {
                   attrs->kernel_layout == "OIHW4o")) {
         supports_texture_storage = true;
       }
+    } else if (auto attrs = call->attrs.as<Conv2DWinogradAttrs>()) {
+      if ((attrs->data_layout == "NCHW4c" || attrs->data_layout == "NHWC4c") &&
+         (attrs->kernel_layout == "OIHW4o" || attrs->kernel_layout == "HWIO4o")) {
+        supports_texture_storage = true;
+      }
     } else if (auto attrs = call->attrs.as<GlobalPool2DAttrs>()) {
       if (attrs->layout == "NCHW4c") {
         supports_texture_storage = true;