Adds asserts

iree-org · Jan 16, 2025 · f4b5379 · f4b5379
1 parent 4a626d6
commit f4b5379
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/compiler/src/iree/compiler/Codegen/Common/TypePropagationPass.cpp b/compiler/src/iree/compiler/Codegen/Common/TypePropagationPass.cpp
@@ -228,6 +228,8 @@ struct GenericOpTypePropagation
       }
       auto inputOperandType =
           llvm::cast<RankedTensorType>(genericOp->getOperandTypes()[index]);
+      assert(inputOperandType.getElementType() == argType &&
+             "expected same element type");
       std::optional<Type> legalizedArgType =
           legalizeStorageElementType(inputOperandType);
       if (!legalizedArgType) {
@@ -259,6 +261,8 @@ struct GenericOpTypePropagation
           modifyYield = true;
           OpOperand *yieldOperand =
               modifiedOp.getMatchingYieldValue(modifiedOpOperand);
+          assert(llvm::cast<TensorType>(modifiedOpOperand->get().getType()).getElementType() ==
+                 yieldOperand->get().getType() && "expected same element type");
           std::optional<Type> legalizedType =
               legalizeStorageElementType(modifiedOpOperand->get().getType());
           if (!legalizedType) {
@@ -289,8 +293,11 @@ struct LinalgFillTypePropagation
   matchAndRewrite(linalg::FillOp fillOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     Value value = adaptor.getInputs().front();
+    TensorType outputType = cast<TensorType>(adaptor.getOutputs()[0].getType());
+    assert(outputType.getElementType() == value.getType() &&
+           "expected same element type");
     std::optional<Type> legalizedElementType =
-        legalizeStorageElementType(adaptor.getOutputs()[0].getType());
+        legalizeStorageElementType(outputType);
     if (!legalizedElementType) {
       return fillOp.emitOpError("failed to get legalized type for value");
     }