moore.lua

local MUSA_ROOT = os.getenv("MUSA_ROOT") or os.getenv("MUSA_HOME") or os.getenv("MUSA_PATH")
add_includedirs(MUSA_ROOT .. "/include")
add_linkdirs(MUSA_ROOT .. "/lib")
add_links("musa", "musart", "mudnn", "mublas")

rule("mu")
    set_extensions(".mu")
    on_load(function (target)
        target:add("includedirs", "include")
    end)

    on_build_file(function (target, sourcefile)
        local objectfile = target:objectfile(sourcefile)
        os.mkdir(path.directory(objectfile))

        local mcc = MUSA_ROOT .. "/bin/mcc"
        local includedirs = table.concat(target:get("includedirs"), " ")

        local args = {
            "-c", sourcefile,
            "-o", objectfile,
            "-I" .. MUSA_ROOT .. "/include",
            "-O3",
            "-fPIC",
            "-Wall",
            "-std=c++17",
            "-pthread"
        }
        local moore_gpu_arch = get_config("moore-gpu-arch")

        if moore_gpu_arch == "mp_31" then
            table.insert(args, 1, "--cuda-gpu-arch=mp_31")
        end

        for _, includedir in ipairs(target:get("includedirs")) do
            table.insert(args, "-I" .. includedir)
        end

        -- ============================
        -- Retrieve all preprocessor defines added to the current target
        local defines = target:get("defines")
        if defines then
            for _, define in ipairs(defines) do
                table.insert(args, "-D" .. define)
            end
        end
        -- ===================================

        os.execv(mcc, args)
        table.insert(target:objectfiles(), objectfile)
    end)
rule_end()

target("infiniop-moore")
    set_kind("static")
    on_install(function (target) end)
    set_languages("cxx17")
    set_warnings("all", "error")
    add_cxflags("-lstdc++", "-fPIC", "-Wno-comment")
    add_cxxflags("-lstdc++", "-fPIC", "-Wno-comment")
    add_files("../src/infiniop/devices/moore/*.cc")
    add_files("../src/infiniop/ops/*/moore/*.mu", {rule = "mu"})

    -- Add source files for Moore muBLAS/muDNN GEMM backends.
    add_files("../src/infiniop/ops/gemm/moore/*/*.mu", {rule = "mu"})

    -- Add source files for Moore per_channel_quant_int8 backends.
    add_files("../src/infiniop/ops/quant/per_channel_quant_int8/moore/*.mu", {rule = "mu"})
target_end()

target("infinirt-moore")
    set_kind("static")
    set_languages("cxx17")
    on_install(function (target) end)
    add_deps("infini-utils")
    set_warnings("all", "error")
    add_cxflags("-lstdc++", "-fPIC")
    add_cxxflags("-lstdc++", "-fPIC")
    add_files("../src/infinirt/moore/*.cc")
target_end()

target("infiniccl-moore")
    set_kind("static")
    add_deps("infinirt")
    on_install(function (target) end)
    set_warnings("all", "error")
    if not is_plat("windows") then
        add_cxflags("-fPIC")
        add_cxxflags("-fPIC")
    end
    if has_config("ccl") then
        add_links("libmccl.so")
        add_files("../src/infiniccl/moore/*.cc")
        
        -- Moore GPU arch with mp_31 support mcclBfloat16 in MCCL
        if get_config("moore-gpu-arch") == "mp_31" then
            add_defines("MARCH_TYPE=310")
            add_cxxflags("-Wno-unused-function")
        end
    end
    set_languages("cxx17")

target_end()