Class SamplingDispatcher

python/tvm/relax/backend/dispatch_sampling.py:30–80

Dispatcher to dispatch sampling op.

@expr_functor.mutator
class SamplingDispatcher(BackendDispatcher):
    """Dispatcher to dispatch sampling op."""

    def visit_call_(self, call: relax.Call) -> relax.Expr:
        if not isinstance(call.op, Op):
            return super().visit_call_(call)

        if call.op.name == "relax.multinomial_from_uniform":
            from tvm.relax.backend.gpu_generic import (  # pylint: disable=import-outside-toplevel
                generic_get_sample_index,
                gpu_multinomial_from_uniform,
            )

            prob, uniform_sample, sample_indices = call.args
            tgt = self._get_target(call.struct_info)
            dtype = call.attrs.dtype
            _, prob_dtype = self.get_shape_dtype(prob)
            sample_shape, sample_dtype = self.get_shape_dtype(uniform_sample)
            sample_indices_shape, sample_indices_dtype = self.get_shape_dtype(sample_indices)

            if len(sample_shape) != 2 or sample_shape[1] != 1:
                raise ValueError("uniform_sample should be a 2D tensor with shape (N, 1)")

            if len(sample_indices_shape) != 2 or sample_indices_shape[1] != 1:
                raise ValueError("sample_indices should be a 2D tensor with shape (N, 1)")

            if self.is_gpu_target(tgt):
                gv = self.builder_.add_func(
                    gpu_multinomial_from_uniform(
                        prob_dtype, sample_dtype, sample_indices_dtype, dtype
                    ),
                    "gpu_multinomial_from_uniform",
                )
                return relax.call_tir(
                    gv,
                    [prob, uniform_sample, sample_indices],
                    out_sinfo=call.struct_info,
                )
            else:
                cumsum_prob = relax.op.cumsum(prob, axis=1, dtype=prob_dtype, exclusive=False)
                gv = self.builder_.add_func(
                    generic_get_sample_index(prob_dtype, sample_dtype, sample_indices_dtype, dtype),
                    "get_sample_index",
                )
                return relax.call_tir(
                    gv,
                    [cumsum_prob, uniform_sample, sample_indices],
                    out_sinfo=call.struct_info,
                )

        return super().visit_call_(call)


@module_pass(opt_level=0, name="DispatchSampling")