[Bug] `relax.transform.SpecializePrimFuncBasedOnCallSite` crashes when a `call_tir` argument is an `R.Prim` scalar

[YuanchengJiang]

The following code:

import tvm
from tvm import relax
from tvm.script import ir as I, relax as R, tirx as T

@I.ir_module
class Module:
    @R.function
    def main(A: R.Tensor((16,), "float16"), scale: R.Prim("float32"), C: R.Tensor((16,), "float16")):
        B = R.call_tir(Module.add_scaled, (A, scale, C), out_ty=R.Tensor((16,), "float16"))
        return B

    @T.prim_func
    def add_scaled(A: T.Buffer((16,), "float16"), scale: T.float32,
                   C: T.Buffer((16,), "float16"), B: T.Buffer((16,), "float16")):
        for i in range(16):
            B[i] = A[i] + T.Cast("float16", scale) * C[i]

relax.analysis.well_formed(Module)  # passes
relax.transform.SpecializePrimFuncBasedOnCallSite()(Module)

Resulted in this output:

Traceback (most recent call last):
  File "min.py", line 19, in <module>
  File "src/relax/transform/specialize_primfunc_based_on_callsite.cc", line 156, in operator()
  File "src/relax/transform/specialize_primfunc_based_on_callsite.cc", line 60, in tvm::relax::SpecializeTIRCallArgs::Run(tvm::IRModule)
  File "src/relax/transform/specialize_primfunc_based_on_callsite.cc", line 74, in tvm::relax::SpecializeTIRCallArgs::VisitExpr_(tvm::CallNode const*)
  File "src/relax/transform/specialize_primfunc_based_on_callsite.cc", line 88, in tvm::relax::Expr tvm::relax::SpecializeTIRCallArgs::SpecializeTirPrimFunc(tvm::Call)
tvm.error.InternalError: Check failed: (ty->IsInstance<TensorTypeNode>()) is false: Expected Tensor struct Info for call :ir.Op(span=None, ty=ir.Type(span=None), name="relax.call_tir", ...)

To reproduce: python3 min.py

Passing an R.Prim value to R.call_tir for a scalar PrimFunc parameter is valid: the module is well-formed, and tvm.compile(Module, target="llvm") builds it without error. SpecializeTIRCallArgs::SpecializeTirPrimFunc (specialize_primfunc_based_on_callsite.cc:86-88) assumes every argument has TensorType and ICHECKs on the PrimType one, instead of skipping non-tensor arguments.

TVM commit: e269315c90e3a061c9e1c77b370ce883b1b223f4

This bug was found by fusion-fuzz