Checking that tvm.relay.expr.Call.op exists

apache · Nov 26, 2021 · c10f7ca · c10f7ca
1 parent d6cad0c
commit c10f7ca
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/python/tvm/relay/backend/contrib/ethosu/codegen.py b/python/tvm/relay/backend/contrib/ethosu/codegen.py
@@ -90,14 +90,14 @@ def visit_call(self, call: tvm.relay.expr.Call) -> tvm.relay.expr.Call:
         new_call = call
         lut_activations = ["TANH", "LUT"]
 
-        if (
-            call.op.name == "contrib.ethosu.identity"
-            and call.attrs.activation in lut_activations
-            and isinstance(call.args[0], tvm.relay.expr.Call)
-        ):
+        if isinstance(call.op, tvm.ir.Op) and isinstance(call.args[0], tvm.relay.expr.Call):
             producer_op = call.args[0]
             # Check if the producer can do a LUT operation
-            if producer_op.op.name in self.lut_ops.keys():
+            if (
+                producer_op.op.name in self.lut_ops.keys()
+                and call.op.name == "contrib.ethosu.identity"
+                and call.attrs.activation in lut_activations
+            ):
                 # Check the producer doesn't already have a LUT
                 has_lut = producer_op.attrs.activation in lut_activations
                 if not has_lut:

diff --git a/python/tvm/relay/backend/contrib/ethosu/op/op_attrs.py b/python/tvm/relay/backend/contrib/ethosu/op/op_attrs.py
@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""The attributes node used for EthosU Relay operators."""
+"""The attributes node used for Arm(R) Ethos(TM)-U NPU Relay operators."""
 from tvm.ir import Attrs
 import tvm._ffi