def lower_to_ir()

in python/featgraph/module/spmm.py [0:0]

• 15 lines of code
• 2 McCabe index (conditional complexity)

    def lower_to_ir(self, input_placeholders, compute_args, schedule_args):
        """Return the IR. This can be useful for debug.

        Parameters
        ----------
        input_placeholders : list of te.placeholder
            The required input tvm placeholders other than adj (which has been passed in during self.init)

        compute_args : dict
            Arguments required for compute_func, e.g., num_feat_partitions

        schedule_args : dict
            Arguments required for schedule_func, e.g., num_cuda_blocks
        """
        if self._num_col_partitions > 1:
            out_placeholder = self._compute_func(*input_placeholders, self._adj_s1_pos_placeholder, \
                self._adj_s1_idx_placeholder, self._adj_vals_placeholder,
                self._adj_d1_size, self._adj_d2_size, **compute_args)
            s = self._schedule_func(out_placeholder, **schedule_args)
            ir = tvm.lower(s, [*input_placeholders, self._adj_s1_pos_placeholder, \
                self._adj_s1_idx_placeholder, self._adj_vals_placeholder, out_placeholder], simple_mode=True)
        else:
            out_placeholder = self._compute_func(*input_placeholders, self._adj_indptr_placeholder, \
                self._adj_indices_placeholder, self._adj_vals_placeholder, **compute_args)
            s = self._schedule_func(out_placeholder, **schedule_args)
            ir = tvm.lower(s, [*input_placeholders, self._adj_indptr_placeholder, \
                self._adj_indices_placeholder, self._adj_vals_placeholder, out_placeholder], simple_mode=True)
        return ir