Method main

tests/python/contrib/test_tir_triton_integration.py:89–94

(x: R.Tensor(("m",), "float32"), y: R.Tensor(("m",), "float32"))

        @R.function
        def main(x: R.Tensor(("m",), "float32"), y: R.Tensor(("m",), "float32")):
            m = T.int64()
            with R.dataflow():
                output = R.call_tir(Module.add, [x, y], relax.TensorStructInfo((m,), "float32"))
                R.output(output)
            return output

    # Constexpr parameters (BLOCK_SIZE) stay in the kernel arguments, and the
    # thread extent is 256 because the kernel is compiled with num_warps=8.