# This is an autogenerated patch file. Do not edit directly. # To update, see instructions at go/patch-triton. diff --git b/BUILD b/BUILD new file mode 100644 --- /dev/null +++ b/BUILD @@ -0,0 +1,1299 @@ +# This package imports OpenAI's Triton (https://github.com/triton-lang/triton). + +# copybara:uncomment load("@rules:copybara.bzl", "copybara_config_test") +load("@rules_cc//cc:cc_binary.bzl", "cc_binary") +load("@rules_cc//cc:cc_library.bzl", "cc_library") +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") +# copybara:uncomment load("//tools/build_defs/license:license.bzl", "license") + +# See go/triton-sla. +package( + # copybara:uncomment_begin + # default_applicable_licenses = [":license"], + # default_compatible_with = [# "@build-target"], + # default_visibility = [ + # # Add your project here if you need to depend on Triton's C++ sources. + # # Add a point of contact we can reach out to when needed in the comment. + # # + # # If you need to use the Python frontend, add your project to + # # //third_party/py/triton/BUILD instead. + # # + # # By adding your project here, you agree to the Triton SLA: go/triton-sla + # "//third_party/py/jax:__subpackages__", # cjfj@ + # "//third_party/tensorflow/compiler/xla:__subpackages__", # bchetioui@ + # "@xla-experimental//gpu:__subpackages__", # csigg@ + # "@xla-experimental//tools/triton:__subpackages__", # vwbaker@ + # "//third_party/py/enzyme_ad:__subpackages__", # wmoses@ + # # Triton-internal visibility + # "@triton//:__subpackages__", + # ], + # copybara:uncomment_end_and_comment_begin + default_visibility = ["//visibility:public"], + # copybara:comment_end + # TODO(csigg): fix and remove + features = [ + "-parse_headers", + "-use_header_modules", + ], +) + +# copybara:uncomment_begin +# license(name = "license") +# +# licenses(["notice"]) +# +# exports_files(["LICENSE"]) +# copybara:uncomment_end + +config_setting( + name = "compiler_is_msvc", + flag_values = { + "@bazel_tools" + # copybara:comment + "//tools/cpp:compiler": "msvc-cl", + }, +) + +# TODO(csigg): fix, enable error upstream, remove. +_no_unused_variable = select({ + ":compiler_is_msvc": [], + "//conditions:default": ["-Wno-unused-variable"], +}) + +# generate version.h for building TritonPluginUtils +genrule( + name = "triton_version_h_gen", + srcs = [ + "include/triton/Version.h.in", + "@xla//third_party/triton:workspace", + ], + outs = ["include/triton/Version.h"], + cmd = "TRITON_COMMIT=$$(sed -n 's/.*TRITON_COMMIT = \"\\(.*\\)\".*/\\1/p' $(location @xla//third_party/triton:workspace)) && " + + "sed \"s/@TRITON_VERSION@/$$TRITON_COMMIT/g\" $(location include/triton/Version.h.in) > $@", +) + +cc_library( + name = "triton_version", + hdrs = ["include/triton/Version.h"], + strip_include_prefix = "include", + deps = [":triton_version_h_gen"], +) + +td_library( + name = "td_files", + srcs = glob(["include/triton/**/*.td"]), + includes = ["include"], + deps = [ + "@llvm-project//mlir:ArithOpsTdFiles", + "@llvm-project//mlir:CastInterfacesTdFiles", + "@llvm-project//mlir:ControlFlowInterfacesTdFiles", + "@llvm-project//mlir:DestinationStyleOpInterfaceTdFiles", + "@llvm-project//mlir:FunctionInterfacesTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", + "@llvm-project//mlir:LLVMOpsTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:PassBaseTdFiles", + "@llvm-project//mlir:SideEffectInterfacesTdFiles", + "@llvm-project//mlir:ViewLikeInterfaceTdFiles", + ], +) + +gentbl_cc_library( + name = "triton_attr_inc_gen", + tbl_outs = { + "include/triton/Dialect/Triton/IR/TritonAttrDefs.h.inc": ["--gen-attrdef-decls"], + "include/triton/Dialect/Triton/IR/TritonAttrDefs.cpp.inc": ["--gen-attrdef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/IR/TritonAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_dialect_inc_gen", + tbl_outs = { + "include/triton/Dialect/Triton/IR/Dialect.h.inc": ["--gen-dialect-decls"], + "include/triton/Dialect/Triton/IR/Dialect.cpp.inc": ["--gen-dialect-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/IR/TritonDialect.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_interfaces_inc_gen", + tbl_outs = { + "include/triton/Dialect/Triton/IR/AttrInterfaces.h.inc": ["--gen-attr-interface-decls"], + "include/triton/Dialect/Triton/IR/AttrInterfaces.cpp.inc": ["--gen-attr-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/IR/TritonInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_ops_inc_gen", + tbl_outs = { + "include/triton/Dialect/Triton/IR/OpsEnums.h.inc": ["--gen-enum-decls"], + "include/triton/Dialect/Triton/IR/OpsEnums.cpp.inc": ["--gen-enum-defs"], + "include/triton/Dialect/Triton/IR/Ops.h.inc": ["--gen-op-decls"], + "include/triton/Dialect/Triton/IR/Ops.cpp.inc": ["--gen-op-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/IR/TritonOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_types_inc_gen", + tbl_outs = { + "include/triton/Dialect/Triton/IR/Types.h.inc": ["--gen-typedef-decls"], + "include/triton/Dialect/Triton/IR/Types.cpp.inc": ["--gen-typedef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/IR/TritonTypes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_transforms_inc_gen", + tbl_outs = {"include/triton/Dialect/Triton/Transforms/Passes.h.inc": [ + "--gen-pass-decls", + "--name=Triton", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/Transforms/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_combine_inc_gen", + # The generated file is #included without relative path. + strip_include_prefix = "lib/Dialect/Triton/Transforms", + tbl_outs = {"lib/Dialect/Triton/Transforms/TritonCombine.inc": ["--gen-rewriters"]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "lib/Dialect/Triton/Transforms/Combine.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_canonicalize_inc_gen", + # The generated file is #included without relative path. + strip_include_prefix = "lib/Dialect/Triton/IR", + tbl_outs = {"lib/Dialect/Triton/IR/TritonCanonicalize.inc": ["--gen-rewriters"]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "lib/Dialect/Triton/IR/Canonicalize.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_attr_interfaces_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/AttrInterfaces.h.inc": ["--gen-attr-interface-decls"], + "include/triton/Dialect/TritonGPU/IR/AttrInterfaces.cpp.inc": ["--gen-attr-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUAttrInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_attr_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/AttrDefs.h.inc": ["--gen-attrdef-decls"], + "include/triton/Dialect/TritonGPU/IR/AttrEnums.h.inc": ["--gen-enum-decls"], + "include/triton/Dialect/TritonGPU/IR/AttrEnums.cpp.inc": ["--gen-enum-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td", + deps = [ + "td_files", + ":triton_gpu_attr_interfaces_inc_gen", + ], +) + +gentbl_cc_library( + name = "triton_gpu_attr_impls_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/AttrDefs.cpp.inc": ["--gen-attrdef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUAttrImpls.td", + deps = [ + "td_files", + ":triton_gpu_attr_inc_gen", + ], +) + +gentbl_cc_library( + name = "triton_gpu_op_interfaces_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/OpInterfaces.h.inc": ["--gen-op-interface-decls"], + "include/triton/Dialect/TritonGPU/IR/OpInterfaces.cpp.inc": ["--gen-op-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUOpInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_cga_encoding_attr_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.h.inc": ["--gen-attrdef-decls"], + "include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.cpp.inc": ["--gen-attrdef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/CGAEncodingAttr.td", + deps = [ + "td_files", + ":triton_gpu_attr_impls_inc_gen", + ], +) + +gentbl_cc_library( + name = "triton_gpu_dialect_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/Dialect.h.inc": ["--gen-dialect-decls"], + "include/triton/Dialect/TritonGPU/IR/Dialect.cpp.inc": ["--gen-dialect-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_instrument_ops_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonInstrument/IR/Ops.h.inc": ["--gen-op-decls"], + "include/triton/Dialect/TritonInstrument/IR/Ops.cpp.inc": ["--gen-op-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonInstrument/IR/TritonInstrumentOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_instrument_attr_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonInstrument/IR/OpsEnums.h.inc": ["--gen-enum-decls"], + "include/triton/Dialect/TritonInstrument/IR/OpsEnums.cpp.inc": ["--gen-enum-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonInstrument/IR/TritonInstrumentAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_enums_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/OpsEnums.h.inc": ["--gen-enum-decls"], + "include/triton/Dialect/TritonGPU/IR/OpsEnums.cpp.inc": ["--gen-enum-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUEnums.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_ops_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/Ops.h.inc": ["--gen-op-decls"], + "include/triton/Dialect/TritonGPU/IR/Ops.cpp.inc": ["--gen-op-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_types_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/Types.h.inc": ["--gen-typedef-decls"], + "include/triton/Dialect/TritonGPU/IR/Types.cpp.inc": ["--gen-typedef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUTypes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_type_interfaces_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonGPU/IR/TypeInterfaces.h.inc": ["--gen-type-interface-decls"], + "include/triton/Dialect/TritonGPU/IR/TypeInterfaces.cpp.inc": ["--gen-type-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/IR/TritonGPUTypeInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_type_interfaces_inc_gen", + tbl_outs = { + "include/triton/Dialect/Triton/IR/TypeInterfaces.h.inc": ["--gen-type-interface-decls"], + "include/triton/Dialect/Triton/IR/TypeInterfaces.cpp.inc": ["--gen-type-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/IR/TritonTypeInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_gpu_transforms_inc_gen", + tbl_outs = {"include/triton/Dialect/TritonGPU/Transforms/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonGPU", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonGPU/Transforms/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_instrument_dialect_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonInstrument/IR/Dialect.h.inc": ["--gen-dialect-decls"], + "include/triton/Dialect/TritonInstrument/IR/Dialect.cpp.inc": ["--gen-dialect-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_instrument_transforms_inc_gen", + tbl_outs = {"include/triton/Dialect/TritonInstrument/Transforms/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonInstrument", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonInstrument/Transforms/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_nvidia_gpu_dialect_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.h.inc": ["--gen-dialect-decls"], + "include/triton/Dialect/TritonNvidiaGPU/IR/Dialect.cpp.inc": ["--gen-dialect-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_nvidia_gpu_ops_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonNvidiaGPU/IR/Ops.h.inc": ["--gen-op-decls"], + "include/triton/Dialect/TritonNvidiaGPU/IR/Ops.cpp.inc": ["--gen-op-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_nvidia_gpu_types_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonNvidiaGPU/IR/Types.h.inc": ["--gen-typedef-decls"], + "include/triton/Dialect/TritonNvidiaGPU/IR/Types.cpp.inc": ["--gen-typedef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUTypes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_nvidia_gpu_op_interfaces_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.h.inc": ["--gen-op-interface-decls"], + "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.cpp.inc": ["--gen-op-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_nvidia_gpu_transforms_inc_gen", + tbl_outs = {"include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonNvidiaGPU", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_conversion_triton_to_triton_gpu_passes_inc_gen", + tbl_outs = {"include/triton/Conversion/TritonToTritonGPU/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonToTritonGPU", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Conversion/TritonToTritonGPU/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_target_llvmir_passes_inc_gen", + tbl_outs = {"include/triton/Target/LLVMIR/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonLLVMIR", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Target/LLVMIR/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_conversion_triton_gpu_to_llvm_pass_inc_gen", + tbl_outs = {"include/triton/Conversion/TritonGPUToLLVM/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonGPUToLLVM", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Conversion/TritonGPUToLLVM/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_op_interfaces_inc_gen", + tbl_outs = { + "include/triton/Dialect/Triton/IR/OpInterfaces.h.inc": ["--gen-op-interface-decls"], + "include/triton/Dialect/Triton/IR/OpInterfaces.cpp.inc": ["--gen-op-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Triton/IR/TritonOpInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "triton_nvidia_gpu_attr_inc_gen", + tbl_outs = { + "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.h.inc": ["--gen-attrdef-decls"], + "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.cpp.inc": ["--gen-attrdef-defs"], + "include/triton/Dialect/TritonNvidiaGPU/IR/OpsEnums.h.inc": ["--gen-enum-decls"], + "include/triton/Dialect/TritonNvidiaGPU/IR/OpsEnums.cpp.inc": ["--gen-enum-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "gluon_attr_inc_gen", + tbl_outs = { + "include/triton/Dialect/Gluon/IR/GluonAttrDefs.h.inc": ["--gen-attrdef-decls"], + "include/triton/Dialect/Gluon/IR/GluonAttrDefs.cpp.inc": ["--gen-attrdef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Gluon/IR/GluonAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "gluon_dialect_inc_gen", + tbl_outs = { + "include/triton/Dialect/Gluon/IR/Dialect.h.inc": ["--gen-dialect-decls"], + "include/triton/Dialect/Gluon/IR/Dialect.cpp.inc": ["--gen-dialect-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Gluon/IR/GluonDialect.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "gluon_ops_inc_gen", + tbl_outs = { + "include/triton/Dialect/Gluon/IR/Ops.h.inc": ["--gen-op-decls"], + "include/triton/Dialect/Gluon/IR/Ops.cpp.inc": ["--gen-op-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Gluon/IR/GluonOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "gluon_transforms_inc_gen", + tbl_outs = {"include/triton/Dialect/Gluon/Transforms/Passes.h.inc": [ + "--gen-pass-decls", + "--name=Gluon", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/triton/Dialect/Gluon/Transforms/Passes.td", + deps = ["td_files"], +) + +cc_library( + name = "TritonDialects", + srcs = glob([ + "lib/Dialect/Gluon/IR/*.cpp", + "lib/Dialect/Triton/IR/*.cpp", + "lib/Dialect/TritonGPU/IR/*.cpp", + "lib/Dialect/TritonInstrument/IR/*.cpp", + "lib/Dialect/TritonNvidiaGPU/IR/*.cpp", + # There are so many interdependencies between Dialect and Analysis that we're just compiling + # everything in a single unit. + "lib/Analysis/*.cpp", + ]) + [ + "include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h", # Avoid circular dependency. + "include/triton/Conversion/TritonGPUToLLVM/Utility.h", # Avoid circular dependency. + "lib/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.cpp", # Avoid circular dependency. + "lib/Dialect/TritonGPU/Transforms/Utility.cpp", # Avoid circular dependency. + "lib/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.cpp", # Avoid circular dependency. + "lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp", # Avoid circular dependency. + ], + hdrs = glob([ + "include/triton/Dialect/Gluon/IR/*.h", + "include/triton/Dialect/Triton/IR/*.h", + "include/triton/Dialect/TritonGPU/IR/*.h", + "include/triton/Dialect/TritonInstrument/IR/*.h", + "include/triton/Dialect/TritonNvidiaGPU/IR/*.h", + # There are so many interdependencies between Dialect and Analysis that we're just compiling + # everything in a single unit. + "include/triton/Analysis/*.h", + ]) + [ + "include/triton/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.h", # Avoid circular dependency. + "include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h", # Avoid circular dependency. + "include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h", # Avoid circular dependency. + "include/triton/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.h", # Avoid circular dependency. + "include/triton/Dialect/TritonGPU/Transforms/Utility.h", # Avoid circular dependency. + # What is this lone header doing rooted under Conversion? Best to add it to Dialect, but + # it would be better if upstream moved it there. + "include/triton/Conversion/MLIRTypes.h", + "include/triton/Dialect/TritonInstrument/Transforms/ConSanTargetHooks.h", + ], + copts = select({ + ":compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + "-Wno-logical-op-parentheses", + "-Wno-unused-but-set-parameter", + "-Wno-ctad-maybe-unsupported", + "-Wno-implicit-fallthrough", + ], + }), + includes = ["include"], + deps = [ + ":Dump", + ":GetEnv", # for Utility + ":TritonTools", + ":gluon_attr_inc_gen", + ":gluon_dialect_inc_gen", + ":gluon_ops_inc_gen", + ":triton_canonicalize_inc_gen", + ":triton_dialect_inc_gen", + ":triton_gpu_attr_impls_inc_gen", + ":triton_gpu_attr_inc_gen", + ":triton_gpu_attr_interfaces_inc_gen", + ":triton_gpu_cga_encoding_attr_inc_gen", + ":triton_gpu_dialect_inc_gen", + ":triton_gpu_enums_inc_gen", + ":triton_gpu_op_interfaces_inc_gen", + ":triton_gpu_ops_inc_gen", + ":triton_gpu_type_interfaces_inc_gen", + ":triton_gpu_types_inc_gen", + ":triton_instrument_attr_inc_gen", + ":triton_instrument_dialect_inc_gen", + ":triton_instrument_ops_inc_gen", + ":triton_interfaces_inc_gen", + ":triton_nvidia_gpu_attr_inc_gen", + ":triton_nvidia_gpu_dialect_inc_gen", + ":triton_nvidia_gpu_op_interfaces_inc_gen", + ":triton_nvidia_gpu_ops_inc_gen", + ":triton_nvidia_gpu_transforms_inc_gen", + ":triton_nvidia_gpu_types_inc_gen", + ":triton_op_interfaces_inc_gen", + ":triton_ops_inc_gen", + ":triton_type_interfaces_inc_gen", + ":triton_types_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:InliningUtils", + "@llvm-project//mlir:LLVMCommonConversion", # for Utility + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:MathDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TensorDialect", + "@llvm-project//mlir:TransformUtils", # for Utility + "@llvm-project//mlir:Transforms", # for Utility + "@llvm-project//mlir:UBDialect", + "@triton//third_party/f2reduce", # for Utility + "@triton//third_party/nvidia:NVGPUDialect", + ], + alwayslink = True, +) + +cc_library( + name = "TritonTransforms", + srcs = glob( + include = [ + "lib/Dialect/Triton/Transforms/*.cpp", + ], + exclude = [ + # Included in TritonGPUTransforms to avoid circular dependency + "lib/Dialect/Triton/Transforms/LoopPeeling.cpp", + ], + ), + hdrs = glob( + include = [ + "include/triton/Dialect/Triton/Transforms/*.h", + ], + exclude = [ + # Included in TritonGPUTransforms to avoid circular dependency + "include/triton/Dialect/Triton/Transforms/LoopPeeling.h", + ], + ), + copts = _no_unused_variable, + deps = [ + ":Dump", + ":GetEnv", + ":TritonDialects", + ":TritonGPUTransforms", + ":triton_combine_inc_gen", + ":triton_transforms_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:FuncTransforms", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SCFTransforms", + "@llvm-project//mlir:SCFUtils", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + ], + alwayslink = True, # TritonDialect uses getCanonicalizationPatterns(). +) + +cc_library( + name = "WarpSpecialization", + srcs = glob( + [ + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/*.cpp", + ], + exclude = [ + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionBuilder.cpp", + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/Partition.cpp", + ], + ), + hdrs = glob( + [ + "include/triton/Dialect/TritonGPU/Transforms/WarpSpecialization/*.h", + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/*.h", + ], + ), + copts = _no_unused_variable, + includes = ["lib/Dialect/TritonGPU/Transforms/WarpSpecialization"], + deps = [ + ":Dump", + ":GetEnv", + ":TritonDialects", + ":TritonGPUTransforms", + ":TritonToTritonGPU", + ":TritonToTritonGPUPasses", + ":triton_gpu_transforms_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:ir_headers", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ArithTransforms", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SCFTransforms", + "@llvm-project//mlir:SCFUtils", + "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TensorDialect", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBDialect", + "@triton//third_party/nvidia:NVWSDialect", + "@triton//third_party/nvidia:NVWSTransforms", + ], +) + +cc_library( + name = "TritonGPUTransforms", + srcs = glob( + include = [ + "lib/Dialect/TritonGPU/Transforms/*.cpp", + "lib/Dialect/TritonGPU/Transforms/Pipeliner/*.cpp", + ], + exclude = [ + "lib/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.cpp", + "lib/Dialect/TritonGPU/Transforms/Utility.cpp", + ], + ) + [ + # TritonTransforms target depends on TritonGPUTransforms. But some files + # in TritonGPUTransforms depend on the headers in TritonTransforms, so + # we need to include them here to avoid circular dependency. + "lib/Dialect/Triton/Transforms/LoopPeeling.cpp", + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionBuilder.cpp", + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/Partition.cpp", + ], + hdrs = glob( + include = [ + "include/triton/Dialect/TritonGPU/Transforms/*.h", + "lib/Dialect/TritonGPU/Transforms/*.h", + "lib/Dialect/TritonGPU/Transforms/Pipeliner/*.h", + ], + exclude = [ + "include/triton/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.h", + "include/triton/Dialect/TritonGPU/Transforms/Utility.h", + ], + ) + [ + # TritonTransforms target depends on TritonGPUTransforms. But some files + # in TritonGPUTransforms depend on the headers in TritonTransforms, so + # we need to include them here to avoid circular dependency. + "include/triton/Dialect/Triton/Transforms/LoopPeeling.h", + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionAttrs.h", + ], + copts = select({ + ":compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-logical-op-parentheses", + "-Wno-reorder-ctor", + "-Wno-return-type", + "-Wno-unused-variable", + "-Wno-string-conversion", + "-Wno-implicit-fallthrough", + ], + }), + includes = [ + "include", + "lib/Dialect/TritonGPU/Transforms", + "lib/Dialect/TritonGPU/Transforms/Pipeliner", + "lib/Dialect/TritonGPU/Transforms/WarpSpecialization", + ], + deps = [ + ":Dump", + ":GetEnv", + ":TritonDialects", + ":TritonGPUToLLVM", + ":TritonToTritonGPUPasses", + ":TritonTools", + ":triton_gpu_transforms_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:ir_headers", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ArithTransforms", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SCFTransforms", + "@llvm-project//mlir:SCFUtils", + "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TensorDialect", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBDialect", + "@triton//third_party/nvidia:NVWSDialectHeader", + ], +) + +cc_library( + name = "TritonGPUToLLVM", + srcs = glob([ + "lib/Conversion/TritonGPUToLLVM/**/*.cpp", + ]), + hdrs = glob([ + "include/triton/Conversion/TritonGPUToLLVM/*.h", + "lib/Conversion/TritonGPUToLLVM/*.h", + ]), + copts = select({ + "//conditions:default": [ + "-Wno-unused-variable", + "-Wno-implicit-fallthrough", + ], + }), + includes = ["include"], + deps = [ + ":TritonDialects", + ":TritonTools", + ":triton_conversion_triton_gpu_to_llvm_pass_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:FuncToLLVM", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:NVVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + ], +) + +# The files in `lib/Conversion/TritonInstrumentToLLVM/` require NVIDIA-specific +# components (e.g. PTXBuilder) from `TritonNVIDIAGPUToLLVM`, which creates +# a circular dependency with `TritonGPUToLLVM`. We avoid the cyclic dependency +# by building these sources together with `TritonNVIDIAGPUToLLVM`. +filegroup( + name = "instrument_to_llvm_srcs", + srcs = glob(["lib/Conversion/TritonInstrumentToLLVM/**/*.cpp"]), + visibility = ["@triton//third_party/nvidia:__subpackages__"], +) + +cc_library( + name = "TritonInstrumentTransforms", + srcs = glob(["lib/Dialect/TritonInstrument/Transforms/*.cpp"]), + hdrs = glob(["include/triton/Dialect/TritonInstrument/Transforms/*.h"]), + copts = select({ + ":compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + deps = [ + ":Dump", + ":GetEnv", + ":TritonDialects", + ":TritonTools", + ":triton_instrument_transforms_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + ], +) + +cc_library( + name = "TritonNvidiaGPUTransforms", + srcs = glob( + [ + "lib/Dialect/TritonNvidiaGPU/Transforms/*.cpp", + ], + # The file included into TritonDialects target which this target depends on. + exclude = [ + "lib/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.cpp", + "lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp", + "lib/Dialect/TritonNvidiaGPU/Transforms/Utility.cpp", + ], + ), + hdrs = glob( + [ + "include/triton/Dialect/TritonNvidiaGPU/Transforms/*.h", + "lib/Dialect/TritonNvidiaGPU/Transforms/*.h", + ], + # The file included into TritonDialects target which this target depends on. + exclude = [ + "include/triton/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.h", + "include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h", + "include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h", + "include/triton/Dialect/TritonNvidiaGPU/Transforms/Utility.h", + ], + ), + copts = select({ + ":compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-ctad-maybe-unsupported", + "-Wno-logical-op-parentheses", + "-Wno-non-virtual-dtor", + "-Wno-return-type", + "-Wno-unused-variable", + "-Wno-private-header", + ], + }), + includes = [ + "include", + "lib/Dialect/TritonNvidiaGPU/Transforms", + ], + deps = [ + ":Dump", + ":GetEnv", + ":TritonDialects", + ":TritonGPUToLLVM", + ":TritonGPUTransforms", + ":TritonInstrumentTransforms", + ":TritonTools", + ":triton_nvidia_gpu_transforms_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBDialect", + ], +) + +cc_library( + name = "TritonToTritonGPUPasses", + hdrs = ["include/triton/Conversion/TritonToTritonGPU/Passes.h"], + deps = [ + ":triton_conversion_triton_to_triton_gpu_passes_inc_gen", + ], +) + +cc_library( + name = "TritonToTritonGPU", + srcs = glob([ + "lib/Conversion/TritonToTritonGPU/*.cpp", + ]), + hdrs = glob( + [ + "include/triton/Conversion/TritonToTritonGPU/*.h", + "lib/Conversion/TritonToTritonGPU/*.h", + ], + exclude = ["include/triton/Conversion/TritonToTritonGPU/Passes.h"], + ), + copts = _no_unused_variable, + includes = ["include"], + deps = [ + ":TritonDialects", + ":TritonGPUTransforms", + ":TritonToTritonGPUPasses", + ":TritonTools", + ":TritonTransforms", + ":triton_conversion_triton_to_triton_gpu_passes_inc_gen", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:UBDialect", + ], +) + +cc_library( + name = "TritonLLVMIR", + srcs = glob([ + "lib/Target/LLVMIR/*.cpp", + ]), + hdrs = glob([ + "include/triton/Target/LLVMIR/*.h", + "lib/Target/LLVMIR/*.h", + ]), + copts = _no_unused_variable, + deps = [ + ":Dump", + ":GetEnv", + ":TritonTransforms", + ":triton_target_llvmir_passes_inc_gen", + "@llvm-project//llvm:Analysis", + "@llvm-project//llvm:BinaryFormat", + "@llvm-project//llvm:Core", + "@llvm-project//llvm:IPO", + "@llvm-project//llvm:IRReader", + "@llvm-project//llvm:InstCombine", + "@llvm-project//llvm:Linker", + "@llvm-project//llvm:MC", + "@llvm-project//llvm:Passes", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:Target", + "@llvm-project//mlir:ArithToLLVM", + "@llvm-project//mlir:BuiltinToLLVMIRTranslation", + "@llvm-project//mlir:ConversionPasses", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:IndexToLLVM", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:LLVMIRTransforms", + "@llvm-project//mlir:LLVMToLLVMIRTranslation", + "@llvm-project//mlir:NVVMToLLVMIRTranslation", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:ROCDLToLLVMIRTranslation", + "@llvm-project//mlir:SCFToControlFlow", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:ToLLVMIRTranslation", + "@llvm-project//mlir:Transforms", + ], +) + +cc_library( + name = "TritonPTX", + srcs = glob([ + "lib/Target/PTX/*.cpp", + ]), + hdrs = glob(["include/triton/Target/PTX/*.h"]), + deps = ["@llvm-project//llvm:Support"], +) + +cc_library( + name = "TritonHSACO", + srcs = glob([ + "lib/Target/HSACO/*.cpp", + ]), + hdrs = glob(["include/triton/Target/HSACO/*.h"]), + deps = [ + ":Dump", + ":GetEnv", + ":TritonLLVMIR", + "@llvm-project//llvm:Core", + "@llvm-project//llvm:ExecutionEngine", + "@llvm-project//llvm:MC", + "@llvm-project//llvm:Scalar", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:Target", + "@llvm-project//llvm:TransformUtils", + "@llvm-project//mlir:ExecutionEngine", + "@llvm-project//mlir:ExecutionEngineUtils", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:LLVMToLLVMIRTranslation", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:ToLLVMIRTranslation", + ], +) + +cc_library( + name = "GetEnv", + hdrs = ["include/triton/Tools/Sys/GetEnv.h"], + strip_include_prefix = "include", +) + +cc_library( + name = "Dump", + hdrs = ["include/triton/Tools/Sys/Dump.h"], + strip_include_prefix = "include", +) + +cc_library( + name = "TritonPluginUtils", + srcs = ["lib/Tools/PluginUtils.cpp"], + hdrs = [ + "include/triton/Tools/PluginUtils.h", + ], + copts = _no_unused_variable, + includes = ["include"], + deps = [ + ":Dump", + ":GetEnv", + ":triton_version", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:PluginsLib", + "@triton//python:ir", + ], +) + +cc_library( + name = "TritonTools", + srcs = glob( + include = ["lib/Tools/*.cpp"], + exclude = ["lib/Tools/PluginUtils.cpp"], + ), + hdrs = glob( + include = ["include/triton/Tools/*.h"], + exclude = ["include/triton/Tools/PluginUtils.h"], + ), + copts = _no_unused_variable, + includes = ["include"], + deps = [ + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@triton//third_party/f2reduce", + ], +) + +cc_library( + name = "GluonTransforms", + srcs = glob(["lib/Dialect/Gluon/Transforms/*.cpp"]), + hdrs = glob(["include/triton/Dialect/Gluon/Transforms/*.h"]), + copts = _no_unused_variable, + deps = [ + ":Dump", + ":GetEnv", + ":TritonDialects", + ":TritonGPUTransforms", + ":TritonTools", + ":gluon_transforms_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + ], + alwayslink = True, +) + +cc_library( + name = "AllPassesAndDialects", + srcs = [ + "include/triton/Conversion/TritonToTritonGPU/Passes.h", + "include/triton/Dialect/TritonInstrument/Transforms/Passes.h", + "include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h", + ], + hdrs = ["bin/RegisterTritonDialects.h"], + includes = ["."], # because it includes third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h + deps = [ + ":GluonTransforms", + ":TritonDialects", + ":TritonGPUToLLVM", + ":TritonGPUTransforms", + ":TritonInstrumentTransforms", + ":TritonLLVMIR", + ":TritonNvidiaGPUTransforms", + ":TritonPluginUtils", + ":TritonToTritonGPU", + ":TritonTransforms", + ":WarpSpecialization", + ":triton_conversion_triton_to_triton_gpu_passes_inc_gen", + ":triton_gpu_cga_encoding_attr_inc_gen", + ":triton_nvidia_gpu_transforms_inc_gen", + "@llvm-project//mlir:AllPassesAndDialects", + "@llvm-project//mlir:RegisterAllPasses", + "@triton//test:ProtonTestTransforms", + "@triton//test:TritonTestAnalysis", + "@triton//test:TritonTestDialect", + "@triton//third_party/amd:TritonAMDGPU", + "@triton//third_party/amd:TritonAMDGPUToLLVM", + "@triton//third_party/amd:TritonAMDGPUTransforms", + "@triton//third_party/nvidia:NVGPUDialect", + "@triton//third_party/nvidia:NVGPUToLLVM", + "@triton//third_party/nvidia:NVHopperTransforms", + "@triton//third_party/nvidia:NVWSDialect", + "@triton//third_party/nvidia:NVWSTransforms", + "@triton//third_party/nvidia:TritonNVIDIAGPUToLLVM", + "@triton//third_party/proton:ProtonGPUToLLVM", + "@triton//third_party/proton:ProtonGPUTransforms", + "@triton//third_party/proton:ProtonIR", + "@triton//third_party/proton:ProtonToProtonGPU", + ], +) + +cc_binary( + name = "triton-opt", + srcs = [ + "bin/triton-opt.cpp", + ], + deps = [ + ":AllPassesAndDialects", + # "@abseil-cpp//absl/base", + "@llvm-project//mlir:MlirOptLib", + "@triton//third_party/amd:TestAMDAnalysis", + "@triton//third_party/proton:ProtonIR", + # "@addr2line", # fixdeps: keep + ], +) + +cc_binary( + name = "triton-llvm-opt", + srcs = [ + "bin/triton-llvm-opt.cpp", + "lib/Target/LLVMIR/LLVMPasses.h", + ], + deps = [ + ":TritonLLVMIR", + # "@abseil-cpp//absl/base", + "@llvm-project//llvm:CodeGen", + "@llvm-project//llvm:Core", + "@llvm-project//llvm:IRReader", + "@llvm-project//llvm:Option", + "@llvm-project//llvm:Passes", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:TargetParser", + # "@addr2line", # fixdeps: keep + ], +) + +# See go/triton-debug for usage. +cc_binary( + name = "triton-reduce", + srcs = ["bin/triton-reduce.cpp"], + deps = [ + ":AllPassesAndDialects", + "@llvm-project//mlir:MlirReduceLib", + "@triton//python:ir", + "@triton//third_party/amd:TestAMDAnalysis", + "@triton//third_party/amd:TritonAMDGPU", + "@triton//third_party/amd:TritonAMDGPUToLLVM", + ], +) + +cc_binary( + name = "triton-tensor-layout", + srcs = ["bin/triton-tensor-layout.cpp"], + deps = [ + ":AllPassesAndDialects", + ":TritonDialects", + # "@abseil-cpp//absl/base", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:AsmParser", + "@llvm-project//mlir:IR", + "@triton//third_party/amd:TestAMDAnalysis", + # "@addr2line", # fixdeps: keep + ], +) + +# copybara:uncomment_begin +# copybara_config_test( +# name = "copybara_config_test", +# config = "copy.bara.sky", +# deps = [ +# # "@copybara/testing:all_bara_sky", +# "@triton//:leakr_badwords.dic", +# "@triton//patches:patch_files", +# "//third_party/xla:copybara_library", +# "@xla//third_party/triton:patch_files", +# ], +# ) +# copybara:uncomment_end + +filegroup( + name = "metadata-file", + srcs = ["METADATA"], +) + +filegroup( + name = "all-files", + srcs = glob(include = ["**/*"]), + data = [ + "@triton//python:all-files", + "@triton//python/test:all-files", + "@triton//python/tutorials:all-files", + "@triton//test:all-files", + "@triton//third_party/amd:all-files", + "@triton//third_party/amd/backend:all-files", + "@triton//third_party/f2reduce:all-files", + "@triton//third_party/nvidia:all-files", + "@triton//third_party/nvidia/backend:all-files", + "@triton//third_party/nvidia/language/cuda:all-files", + "@triton//third_party/proton:all-files", + "@triton//third_party/proton/proton:all-files", + "@triton//third_party/proton/test:all-files", + "@triton//unittest:all-files", + ], +) diff --git b/python/BUILD b/python/BUILD new file mode 100644 --- /dev/null +++ b/python/BUILD @@ -0,0 +1,210 @@ +# NOTE: Do not depend on any targets from this directory, +# but use //third_party/py/triton instead. + +load("@rules_cc//cc:cc_library.bzl", "cc_library") +load("@xla//third_party/rules_python/python:py_extension.bzl", "py_extension") +load("@local_config_cuda//cuda:build_defs.bzl", "cuda_library") +load("@pybind11_bazel//:build_defs.bzl", "pybind_extension") + +package( + # copybara:uncomment_begin + # default_applicable_licenses = ["@triton//:license"], + # default_compatible_with = [ + # # "@build-target", + # ], + # default_visibility = [ + # "//third_party/py/triton:__pkg__", + # "@triton//python:__subpackages__", + # ], + # copybara:uncomment_end_and_comment_begin + default_visibility = ["//visibility:public"], + # copybara:comment_end +) + +cc_library( + name = "passes", + hdrs = ["src/passes.h"], + includes = ["src"], + visibility = ["@triton//third_party:__subpackages__"], + deps = [ + "@triton//:Dump", + "@triton//:GetEnv", + ], +) + +cc_library( + name = "ir", + hdrs = ["src/ir.h"], + copts = ["-Wno-unused-variable"], + includes = ["src"], + visibility = ["@triton//:__subpackages__"], + deps = [ + "@llvm-project//mlir:IR", + "@triton//:Dump", + "@triton//:GetEnv", + "@triton//:TritonGPUToLLVM", + ], +) + +cc_library( + name = "gsan", + srcs = ["triton/experimental/gsan/src/gsan_testing.cc"], + hdrs = ["triton/experimental/gsan/src/GSan.h"], + copts = ["-fexceptions"], + features = ["-use_header_modules"], + includes = [".."], + visibility = ["@triton//:__subpackages__"], + deps = [ + "//third_party/nanobind", + "//third_party/python_runtime:headers", + ], +) + +pybind_extension( + name = "gsan_allocator", + srcs = [ + "triton/experimental/gsan/src/GSan.h", + "triton/experimental/gsan/src/GSanAllocator.cc", + ], + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + deps = [ + # "@libcuda", + "@local_config_cuda//cuda:cuda_headers", + "@local_config_cuda//cuda:cuda_runtime", + ], +) + +cuda_library( + name = "gsan_ll_lib", + srcs = ["triton/experimental/gsan/src/GSanLibrary.cu"], + hdrs = [ + "triton/experimental/gsan/src/GSan.h", + "triton/experimental/gsan/src/Hash.cuh", + ], + compatible_with = [ + # Keeping in a separate line for copybara to replace. + # "@build-target", + ], + copts = [ + "-emit-llvm", + "--cuda-device-only", + "--no-cuda-gpu-arch=all", + "--cuda-gpu-arch=sm_80", + "-fcuda-flush-denormals-to-zero", + ], + linkstatic = 1, +) + +genrule( + name = "gsan_ll", + srcs = [":gsan_ll_lib"], + outs = ["gsan.ll"], + cmd = "$(AR) xo $$(echo $(SRCS) | tr ' ' '\\n' | grep libgsan_ll_lib.a) && " + + "$(location @llvm-project//llvm:llvm-dis) GSanLibrary.*o -o $(OUTS)", + toolchains = ["//tools/cpp:toolchain_type"], + tools = ["@llvm-project//llvm:llvm-dis"], + visibility = [ + "//third_party/py/triton:__subpackages__", + "@triton//:__subpackages__", + ], +) + +py_extension( + name = "libtriton", + srcs = [ + "src/gluon_ir.cc", + "src/interpreter.cc", + "src/ir.cc", + "src/linear_layout.cc", + "src/llvm.cc", + "src/main.cc", + "src/passes.cc", + "src/specialize.cc", + ], + copts = [ + "-fexceptions", + "-DTRITON_BACKENDS_TUPLE=(nvidia)", + "-Wno-string-conversion", + ], + features = ["-use_header_modules"], + deps = [ + ":gsan", + ":ir", + ":passes", + "@llvm-project//llvm:Analysis", + "@llvm-project//llvm:CodeGen", + "@llvm-project//llvm:Core", + "@llvm-project//llvm:IPO", + "@llvm-project//llvm:IRReader", + "@llvm-project//llvm:InstCombine", + "@llvm-project//llvm:Instrumentation", + "@llvm-project//llvm:Linker", + "@llvm-project//llvm:MC", + "@llvm-project//llvm:Passes", + "@llvm-project//llvm:Plugins", + "@llvm-project//llvm:Scalar", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:Target", + "@llvm-project//mlir:BuiltinToLLVMIRTranslation", + "@llvm-project//mlir:BytecodeWriter", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:ConversionPasses", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:IndexDialect", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:LLVMIRTransforms", + "@llvm-project//mlir:LLVMToLLVMIRTranslation", + "@llvm-project//mlir:NVVMToLLVMIRTranslation", + "@llvm-project//mlir:Parser", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:ROCDLDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:ToLLVMIRTranslation", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBDialect", + "//third_party/nanobind", + "//third_party/python_runtime:headers", + "@triton//:AllPassesAndDialects", + "@triton//:Dump", + "@triton//:GetEnv", + "@triton//:GluonTransforms", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + "@triton//:TritonGPUTransforms", + "@triton//:TritonHSACO", + "@triton//:TritonInstrumentTransforms", + "@triton//:TritonLLVMIR", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//:TritonPTX", + "@triton//:TritonPluginUtils", + "@triton//:TritonToTritonGPU", + "@triton//:TritonToTritonGPUPasses", + "@triton//:TritonTools", + "@triton//:TritonTransforms", + "@triton//third_party/amd:TDMUtility", + "@triton//third_party/amd:TritonAMDGPU", + "@triton//third_party/amd:TritonAMDGPUTransforms", + "@triton//third_party/nvidia:triton_nvidia", + "@triton//third_party/proton:ProtonIR", + ], +) + +filegroup( + name = "py_files", + srcs = glob( + include = ["triton/**/*.py"], + ), +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/python/test/BUILD b/python/test/BUILD new file mode 100644 --- /dev/null +++ b/python/test/BUILD @@ -0,0 +1,434 @@ +load("//third_party/py/pytest:pytest_defs.bzl", "pytest_multi_tests", "pytest_test") + +package( + default_applicable_licenses = ["@triton//:license"], +) + +_requires_gpu_sm80 = [ + "config-cuda-only", + "requires-gpu-sm80", +] + +_requires_gpu_sm90 = [ + "config-cuda-only", + "requires-gpu-sm90", +] + +_requires_gpu_sm100 = [ + "config-cuda-only", + "requires-gpu-sm100-only", +] + +_requires_gpu_sm100_full = [ + "config-cuda-only", + "requires-gpu-sm100-full-only", +] + +_requires_config_cuda = select( + {"@rules_cuda//cuda:is_enabled": []}, + no_match_error = "Requires --config=cuda", +) + +_asan_needs_extra_link_memory = select({ + "//tools/cpp:asan_build": {"cpp_link.mem": "20g"}, + "//conditions:default": None, +}) + +DEVICELESS_TARGETS = [ + # go/keep-sorted start + "unit/language/test_compile_only.py", + "unit/language/test_frontend.py", + "unit/language/test_mxfp.py", + "unit/plugins/test_plugin.py", + "unit/runtime/test_specialize.py", + "unit/test_filecheck.py", + "unit/tools/test_linear_layout.py", + "unit/tools/test_slice_kernel.py", + "unit/tools/test_stable_toposort.py", + # go/keep-sorted end +] + +EXCLUDE_TESTS = [ + # failing tests + "unit/language/test_subprocess.py", # TODO(b/320224484): fix failing test + "unit/tools/test_aot.py", # TODO(b/320224484): fix failing test + "unit/tools/test_disasm.py", # TODO(b/320224484): fix failing test + "unit/runtime/test_blaslt.py", # TODO(b/346755023): fix failing test + "unit/test_debug.py", # TODO(b/374733875): fix failing test. Also see b/374733872. + # As of Feb 2026, there is no plan to support Gluon in g3. There is no harm in running them, + # but fixing failures is not a priority, so it's OK to disable them. + "gluon/test_consan.py", + "gluon/test_core.py", # Disabled because of b/486813140, could be PTXAS bug. + "gluon/test_frontend.py", + "gluon/test_lowerings.py", + + # valueless internally + "backend/test_device_backend.py", # we run backends differently in g3 + "unit/language/test_reproducer.py", # this is not an actual test, but a tool for running reproducers + "unit/test_knobs.py", # testing binary and file path utilites, which do not work in google's environment + "unit/test_perf_warning.py", # No backtraces in non-debug builds. + "unit/runtime/test_build.py", # this tests OSS launcher, we implement this differently in g3 + "unit/runtime/test_no_torch_dispatch.py", # we run tests differently in g3 & we only use torch in g3 for testing, see third_party/py/triton/google/BUILD + + # requires 2 GPUs to run + "unit/runtime/test_peer_access.py", # this requires 2 GPUs to run, we have a separate target for it + + # sharded tests + "unit/language/test_core.py", # this test is huge and we shard it separately + "unit/runtime/test_driver.py", # Some of the test cases does not use the + # device and thus suite fails with go/forge-accel-or-fail. We run it as a + # single shard to avoid this. + "unit/runtime/test_cache_determinism.py", # Only 2 test cases, so sharding causes go/forge-accel-or-fail on empty shards. We run it as a single shard. +] + DEVICELESS_TARGETS # targets that don't require a GPU to run + +B200_FULL_GPU_ONLY_TESTS = [ + "unit/language/test_matmul.py", # TODO(b/462585714):this test fails when run on B200 MIGs + "gsan/test_gsan.py", # TODO(b/495776594): these tests complain about misaligned access on B200. + "gsan/test_gsan_failures.py", +] + +test_suite( + name = "a100_tests", + tests = [ + ":ampere", + ":unit/language/test_core_a100", + ":unit/runtime/test_cache_determinism_a100", + ":unit/runtime/test_driver_a100", + ], +) + +test_suite( + name = "h100_tests", + tests = [ + ":hopper", + ":unit/language/test_core_h100", + ":unit/runtime/test_cache_determinism_h100", + ":unit/runtime/test_driver_h100", + ":unit/runtime/test_peer_access_h100_x2", + ], +) + +test_suite( + name = "b200_tests", + tests = [ + ":blackwell", + ":blackwell_full", + ":unit/language/test_core_b200", + ":unit/runtime/test_cache_determinism_b200", + ":unit/runtime/test_driver_b200", + ":unit/runtime/test_peer_access_b200_x2", + ], +) + +# Runs all python tests on A100. +pytest_multi_tests( + name = "ampere", + size = "large", + srcs = [ + "conftest.py", + "unit/language/test_core.py", + "unit/language/test_mxfp.py", + ], + exec_properties = _asan_needs_extra_link_memory, + g3_multiprocessing = True, + name_suffix = "_a100", + shard_count = 10, + tags = _requires_gpu_sm80, + target_compatible_with = _requires_config_cuda, + tests = glob( + include = ["**/test_*.py"], + exclude = EXCLUDE_TESTS, + ), + deps = [ + "@llvm-project//llvm:FileCheck", + "//third_party/py/expecttest", + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +# Runs all python tests on H100. +pytest_multi_tests( + name = "hopper", + size = "large", + srcs = [ + "conftest.py", + "unit/language/test_core.py", + "unit/language/test_mxfp.py", + ], + exec_properties = _asan_needs_extra_link_memory, + g3_multiprocessing = True, + name_suffix = "_h100", + shard_count = 10, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + tests = glob( + include = ["**/test_*.py"], + exclude = EXCLUDE_TESTS, + ), + deps = [ + "@llvm-project//llvm:FileCheck", + "//third_party/py/expecttest", + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +# Shard test_core more on H100, as it is otherwise very slow to run. +pytest_test( + name = "unit/language/test_core_h100", + size = "large", + srcs = [ + "conftest.py", + ], + exec_properties = _asan_needs_extra_link_memory, + shard_count = 40, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + tests = ["unit/language/test_core.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +# Shard test_core more on A100, as it is otherwise very slow to run. +pytest_test( + name = "unit/language/test_core_a100", + size = "large", + srcs = [ + "conftest.py", + ], + exec_properties = _asan_needs_extra_link_memory, + shard_count = 40, + tags = _requires_gpu_sm80, + target_compatible_with = _requires_config_cuda, + tests = ["unit/language/test_core.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_test( + name = "unit/runtime/test_driver_a100", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + # Some of the test cases does not use the device and thus suite fails with go/forge-accel-or-fail. + # We run it as a single shard to avoid this. + shard_count = 1, + tags = _requires_gpu_sm80, + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_driver.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_test( + name = "unit/runtime/test_driver_h100", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + # Some of the test cases does not use the device and thus suite fails with go/forge-accel-or-fail. + # We run it as a single shard to avoid this. + shard_count = 1, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_driver.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_test( + name = "unit/runtime/test_driver_b200", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + # Some of the test cases does not use the device and thus suite fails with go/forge-accel-or-fail. + # We run it as a single shard to avoid this. + shard_count = 1, + tags = _requires_gpu_sm100, + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_driver.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_test( + name = "unit/runtime/test_cache_determinism_a100", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + shard_count = 1, + tags = _requires_gpu_sm80, + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_cache_determinism.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_test( + name = "unit/runtime/test_cache_determinism_h100", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + shard_count = 1, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_cache_determinism.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_test( + name = "unit/runtime/test_cache_determinism_b200", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + shard_count = 1, + tags = _requires_gpu_sm100, + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_cache_determinism.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_test( + name = "deviceless", + srcs = [ + "conftest.py", + "unit/plugins/custom_stages.py", + ], + exec_properties = _asan_needs_extra_link_memory, + tests = DEVICELESS_TARGETS, + deps = [ + "@llvm-project//llvm:FileCheck", + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +# Runs all python tests on B200. +pytest_multi_tests( + name = "blackwell", + size = "large", + srcs = [ + "conftest.py", + "unit/language/test_core.py", + "unit/language/test_mxfp.py", + ], + exec_properties = _asan_needs_extra_link_memory, + g3_multiprocessing = True, + name_suffix = "_b200", + shard_count = 10, + tags = _requires_gpu_sm100, + target_compatible_with = _requires_config_cuda, + tests = glob( + include = ["**/test_*.py"], + exclude = EXCLUDE_TESTS + B200_FULL_GPU_ONLY_TESTS, + ), + deps = [ + "@llvm-project//llvm:FileCheck", + "//third_party/py/expecttest", + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +pytest_multi_tests( + name = "blackwell_full", + size = "large", + srcs = [ + "conftest.py", + "unit/language/test_core.py", + "unit/language/test_mxfp.py", + ], + exec_properties = _asan_needs_extra_link_memory, + g3_multiprocessing = True, + name_suffix = "_b200", + shard_count = 10, + tags = _requires_gpu_sm100_full, + target_compatible_with = _requires_config_cuda, + tests = glob( + include = B200_FULL_GPU_ONLY_TESTS, + ), + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +# Shard test_core more on B200, as it is otherwise very slow to run. +pytest_test( + name = "unit/language/test_core_b200", + size = "large", + srcs = [ + "conftest.py", + ], + exec_properties = _asan_needs_extra_link_memory, + shard_count = 40, + tags = _requires_gpu_sm100, + target_compatible_with = _requires_config_cuda, + tests = ["unit/language/test_core.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +# Requires 2 GPUs to run. +pytest_test( + name = "unit/runtime/test_peer_access_h100_x2", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + tags = [ + "config-cuda-only", + "requires-gpu-sm90-full:2", + ], + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_peer_access.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +# Requires 2 GPUs to run. +pytest_test( + name = "unit/runtime/test_peer_access_b200_x2", + size = "large", + srcs = ["conftest.py"], + exec_properties = _asan_needs_extra_link_memory, + tags = [ + "config-cuda-only", + "requires-gpu-sm100-full:2", + ], + target_compatible_with = _requires_config_cuda, + tests = ["unit/runtime/test_peer_access.py"], + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/python/tutorials/BUILD b/python/tutorials/BUILD new file mode 100644 --- /dev/null +++ b/python/tutorials/BUILD @@ -0,0 +1,193 @@ +load("@rules_python//python:defs.bzl", "py_library") +load("@rules_python//python:defs.bzl", "py_test") + +package( + default_applicable_licenses = ["@triton//:license"], +) + +_requires_gpu_sm90 = [ + "config-cuda-only", + "requires-gpu-sm90", +] + +_requires_gpu_sm100 = [ + "config-cuda-only", + "requires-gpu-sm100-full", + "notap", # Blackwell support is experimental on Forge and requires special flags. + "manual", # Blackwell support is experimental on Forge and requires special flags. +] + +_requires_config_cuda = select( + {"@rules_cuda//cuda:is_enabled": []}, + no_match_error = "Requires --config=cuda", +) + +# Package some of these tests that can be used as modules into a library, so we can use the kernel +# implementations elsewhere. +py_library( + name = "kernels", + testonly = True, # Has pytest as a dependency + srcs = [ + "06-fused-attention.py", + "09-persistent-matmul.py", + "10-block-scaled-matmul.py", + ], + strict_deps = False, + visibility = ["//third_party/py/triton:__subpackages__"], + deps = [ + "//third_party/py/pytest", + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + "//third_party/py/triton/profiler", + ], +) + +py_test( + name = "01-vector-add", + srcs = ["01-vector-add.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "02-fused-softmax", + srcs = ["02-fused-softmax.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "03-matrix-multiplication", + srcs = ["03-matrix-multiplication.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "04-low-memory-dropout", + srcs = ["04-low-memory-dropout.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "05-layer-norm", + srcs = ["05-layer-norm.py"], + strict_deps = False, + # Test is 2% flaky b/497760029 + tags = _requires_gpu_sm90 + ["notap"], + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "06-fused-attention", + timeout = "long", # 900s + srcs = ["06-fused-attention.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/pytest", + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "07-extern-functions", + srcs = ["07-extern-functions.py"], + strict_deps = False, + tags = _requires_gpu_sm90 + [ + # Can't find libdevice.10.bc + "manual", + "notap", + ], + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "08-grouped-gemm", + srcs = ["08-grouped-gemm.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +py_test( + name = "09-persistent-matmul", + timeout = "long", # 900s + srcs = ["09-persistent-matmul.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + "//third_party/py/triton/profiler", + ], +) + +py_test( + name = "10-block-scaled-matmul", + srcs = ["10-block-scaled-matmul.py"], + strict_deps = False, + tags = _requires_gpu_sm100, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + "//third_party/py/triton/profiler", + ], +) + +py_test( + name = "11-programmatic-dependent-launch", + srcs = ["11-programmatic-dependent-launch.py"], + strict_deps = False, + tags = _requires_gpu_sm90, + target_compatible_with = _requires_config_cuda, + deps = [ + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/test/BUILD b/test/BUILD new file mode 100644 --- /dev/null +++ b/test/BUILD @@ -0,0 +1,119 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") + +# copybara:uncomment_begin +# load("//third_party/llvm/build_defs:lit.bzl", "glob_lit_tests") +# load("//tools/build_defs/build_test:build_test.bzl", "build_test") +# +# package( +# default_applicable_licenses = ["@triton//:license"], +# default_compatible_with = [# "@build-target"], +# default_visibility = ["@triton//:__subpackages__"], +# ) +# +# glob_lit_tests( +# name = "all_tests", +# data = [ +# "@llvm-project//llvm:FileCheck", +# "@llvm-project//llvm:llc", +# "@llvm-project//llvm:not", +# "@llvm-project//llvm:opt", +# "@llvm-project//llvm:split-file", +# "@llvm-project//mlir:mlir-translate", +# "@triton//:triton-llvm-opt", +# "@triton//:triton-opt", +# "@triton//:triton-tensor-layout", +# ], +# driver = "@llvm-project//mlir:run_lit.sh", +# exclude = [ +# "Conversion/amd/dedup-by-constancy.mlir", # AMD-specific, broken +# "Conversion/amd/wmma-v2-shortcut.mlir", # AMD-specific, broken on it's first integeration. +# "TritonGPU/amd/amd-instruction-sched.mlir", # AMD-specific, broken with -debug-only. +# "TritonGPU/optimize_epilogue.mlir", # TODO: b/346283526 - AMD-specific, triggering UBSAN +# "Conversion/allocate_warp_groups.mlir", # warpGroupStartIds order is wrong, b/418732604 +# # TODO(b/478678840): Re-enable when https://github.com/triton-lang/triton/commit/6559dea9d66c20590b4a7aa0cafa850b55269a1e lands. +# "TritonGPU/amd/amd-range-analysis.mlir", +# "TritonGPU/amd/amd-schedule-hint.mlir", +# ], +# test_file_exts = [ +# "mlir", +# "ll", +# ], +# ) +# +# build_test( +# name = "build_test", +# allow_empty_target = True, +# targets = [ +# "@triton//:TritonDialects", +# "@triton//:TritonGPUToLLVM", +# "@triton//:TritonGPUTransforms", +# "@triton//:TritonLLVMIR", +# "@triton//:TritonPTX", +# "@triton//:TritonToTritonGPU", +# "@triton//:TritonTools", +# "@triton//:TritonTransforms", +# "@triton//:triton-opt", +# ], +# ) +# copybara:uncomment_end + +cc_library( + name = "TritonTestAnalysis", + srcs = glob(["lib/Analysis/*.cpp"]), + hdrs = glob(["include/Analysis/*.h"]), + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-private-header", + ], + }), + visibility = ["//visibility:public"], + deps = [ + "@llvm-project//llvm:Support", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFToControlFlow", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@triton//:TritonDialects", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//third_party/nvidia:TritonNVIDIAGPUToLLVM", + ], +) + +cc_library( + name = "TritonTestDialect", + srcs = glob(["lib/Dialect/*.cpp"]), + visibility = ["//visibility:public"], + deps = [ + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFToControlFlow", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@triton//:TritonDialects", + "@triton//:TritonGPUTransforms", + "@triton//:TritonTransforms", + "@triton//third_party/nvidia:TritonNVIDIAGPUToLLVM", + ], +) + +cc_library( + name = "ProtonTestTransforms", + srcs = ["lib/Proton/TestScopeIdAllocation.cpp"], + visibility = ["//visibility:public"], + deps = [ + "@llvm-project//mlir:Pass", + "@triton//third_party/proton:ProtonAnalysis", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/amd/BUILD b/third_party/amd/BUILD new file mode 100644 --- /dev/null +++ b/third_party/amd/BUILD @@ -0,0 +1,383 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") + +package( + # copybara:uncomment_begin + # default_applicable_licenses = ["@triton//:license"], + # default_compatible_with = [# "@build-target"], + # default_visibility = [ + # "//third_party/py/enzyme_ad:__subpackages__", + # "//third_party/tensorflow/compiler/xla/backends/gpu/codegen/triton:__subpackages__", + # "//third_party/tensorflow/compiler/xla/backends/gpu/tests:__subpackages__", + # "@triton//:__subpackages__", + # ], + # copybara:uncomment_end_and_comment_begin + default_visibility = ["//visibility:public"], + # copybara:comment_end +) + +# TODO(csigg): fix, enable error upstream, remove. +_no_unused_variable = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": ["-Wno-unused-variable"], +}) + +cc_library( + name = "TritonAMDGPUTransforms", + srcs = glob( + [ + "lib/TritonAMDGPUTransforms/**/*.h", + "lib/TritonAMDGPUTransforms/**/*.cpp", + ], + exclude = [ + "lib/TritonAMDGPUTransforms/Utility.h", + "lib/TritonAMDGPUTransforms/MfmaGroup.cpp", # Avoid circular dependency. + "lib/TritonAMDGPUTransforms/WmmaGroup.cpp", # Avoid circular dependency. + ], + ), + hdrs = glob([ + "include/TritonAMDGPUTransforms/**/*.h", + ]) + ["lib/TritonAMDGPUTransforms/Utility.h"], + copts = _no_unused_variable + select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-implicit-fallthrough", + "-Wno-non-virtual-dtor", + ], + }), + includes = [ + "include", + ], + deps = [ + ":Analysis", + ":TDMUtility", + ":TritonAMDGPU", + ":TritonAMDGPUToLLVM", + ":triton_conversion_amdgpu_transforms_passes_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:TargetParser", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:ConvertToLLVM", + "@llvm-project//mlir:FuncTransforms", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:ROCDLDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TensorDialect", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBDialect", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + "@triton//:TritonGPUTransforms", + "@triton//:TritonTools", + ], +) + +cc_library( + name = "TDMUtility", + srcs = glob( + [ + "lib/TritonAMDGPUToLLVM/TDMUtility.*", + ], + ) + [ + "lib/Dialect/TritonAMDGPU/IR/TargetFeatures.cpp", + ], + hdrs = [ + "include/Dialect/TritonAMDGPU/IR/TargetFeatures.h", + "lib/TritonAMDGPUToLLVM/TDMUtility.h", + "lib/TritonAMDGPUToLLVM/TargetInfo.h", + ], + copts = _no_unused_variable, + includes = [ + "include", + "lib/TritonAMDGPUToLLVM", + ], + deps = [ + "@llvm-project//llvm:Support", + "@llvm-project//llvm:TargetParser", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:ROCDLDialect", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + "@triton//:TritonTools", + "@triton//third_party/amd/backend:headers", + ], +) + +cc_library( + name = "TritonAMDGPU", + srcs = glob( + [ + "lib/Dialect/TritonAMDGPU/**/*.h", + "lib/Dialect/TritonAMDGPU/**/*.cpp", + ], + exclude = ["lib/Dialect/TritonAMDGPU/IR/TargetFeatures.cpp"], + ), + hdrs = glob([ + "include/Dialect/TritonAMDGPU/**/*.h", + "include/Utils/**/*.h", + ]), + copts = _no_unused_variable, + includes = [ + "..", + "include", + ], + deps = [ + ":TDMUtility", + ":triton_amdgpu_attr_def_inc_gen", + ":triton_amdgpu_dialect_inc_gen", + ":triton_amdgpu_op_interfaces_inc_gen", + ":triton_amdgpu_ops_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TensorDialect", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + "@triton//:TritonTools", + ], +) + +cc_library( + name = "TritonAMDGPUToLLVM", + srcs = glob( + [ + "lib/TritonAMDGPUToLLVM/**/*.h", + "lib/TritonAMDGPUToLLVM/**/*.cpp", + # TritonAMDGPUToLLVM and TritonAMDGPUDialectToLLVM have interdependencies, easiest way to + # deal with circular dependencies is to just compile both in a single unit. + "lib/TritonAMDGPUDialectToLLVM/**/*.h", + "lib/TritonAMDGPUDialectToLLVM/**/*.cpp", + ], + exclude = [ + "lib/TritonAMDGPUToLLVM/AsyncUtility.h", + "lib/TritonAMDGPUToLLVM/PatternTritonGPUOpToLLVM.h", + "lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.*", + "lib/TritonAMDGPUToLLVM/OptimizeLDSUsage.*", + # TDMUtility is also used by both TritonAMDGPU and TritonAMDGPUToLLVM that would cause + # circular dependencies. Moved TDMUtility to its own target. + "lib/TritonAMDGPUToLLVM/TDMUtility.*", + "lib/TritonAMDGPUToLLVM/SchedInstructions.cpp", + ], + ) + [ + "include/TritonAMDGPUTransforms/MfmaGroup.h", # Avoid circular dependency. + "include/TritonAMDGPUTransforms/WmmaGroup.h", # Avoid circular dependency. + "lib/TritonAMDGPUTransforms/MfmaGroup.cpp", # Avoid circular dependency. + "lib/TritonAMDGPUTransforms/WmmaGroup.cpp", # Avoid circular dependency. + ], + hdrs = glob([ + "include/TritonAMDGPUToLLVM/**/*.h", + ]) + [ + "lib/TritonAMDGPUToLLVM/AsyncUtility.h", + "lib/TritonAMDGPUToLLVM/PatternTritonGPUOpToLLVM.h", + "lib/TritonAMDGPUToLLVM/TargetInfo.h", + "lib/TritonAMDGPUToLLVM/Utility.h", + ], + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + "-Wno-implicit-fallthrough", + "-Wno-sometimes-uninitialized", + "-Wno-reorder-ctor", + ], + }), + includes = [ + "include", + "lib/TritonAMDGPUToLLVM", + ], + deps = [ + ":Analysis", + ":TDMUtility", + ":TritonAMDGPU", + ":triton_conversion_amdgpu_to_llvm_passes_inc_gen", + "@llvm-project//llvm:Passes", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:TargetParser", + "@llvm-project//llvm:ir_headers", + "@llvm-project//mlir:AMDGPUDialect", + "@llvm-project//mlir:AMDGPUTransforms", + "@llvm-project//mlir:AMDGPUUtils", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ArithToLLVM", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:ControlFlowToLLVM", + "@llvm-project//mlir:ConversionPasses", + "@llvm-project//mlir:ConvertToLLVM", + "@llvm-project//mlir:DialectUtils", + "@llvm-project//mlir:FuncToLLVM", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:GPUToNVVMTransforms", + "@llvm-project//mlir:GPUToROCDLTransforms", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:IndexDialect", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:MathToLLVM", + "@llvm-project//mlir:NVVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:ROCDLDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SCFToControlFlow", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBToLLVM", + "@triton//:Dump", + "@triton//:GetEnv", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + "@triton//:TritonTools", + ], +) + +td_library( + name = "td_files", + srcs = glob(["include/**/*.td"]), + includes = ["include"], + deps = ["@triton//:td_files"], +) + +gentbl_cc_library( + name = "triton_amdgpu_ops_inc_gen", + tbl_outs = { + "include/Dialect/TritonAMDGPU/IR/OpsConversions.inc": [ + "--gen-llvmir-conversions", + ], + "include/Dialect/TritonAMDGPU/IR/Ops.h.inc": [ + "--gen-op-decls", + ], + "include/Dialect/TritonAMDGPU/IR/Ops.cpp.inc": [ + "--gen-op-defs", + ], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td", + deps = [":td_files"], +) + +gentbl_cc_library( + name = "triton_amdgpu_op_interfaces_inc_gen", + tbl_outs = { + "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOpInterfaces.h.inc": ["--gen-op-interface-decls"], + "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOpInterfaces.cpp.inc": ["--gen-op-interface-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOpInterfaces.td", + deps = [":td_files"], +) + +gentbl_cc_library( + name = "triton_amdgpu_dialect_inc_gen", + tbl_outs = { + "include/Dialect/TritonAMDGPU/IR/Dialect.h.inc": [ + "--gen-dialect-decls", + "--dialect=amdg", + ], + "include/Dialect/TritonAMDGPU/IR/Dialect.cpp.inc": [ + "--gen-dialect-defs", + "--dialect=amdg", + ], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td", + deps = [":td_files"], +) + +gentbl_cc_library( + name = "triton_amdgpu_attr_def_inc_gen", + tbl_outs = { + "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.h.inc": ["--gen-attrdef-decls"], + "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.cpp.inc": ["--gen-attrdef-defs"], + "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUEnums.h.inc": ["--gen-enum-decls"], + "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUEnums.cpp.inc": ["--gen-enum-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.td", + deps = [":td_files"], +) + +gentbl_cc_library( + name = "triton_conversion_amdgpu_to_llvm_passes_inc_gen", + tbl_outs = {"include/TritonAMDGPUToLLVM/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonAMDGPUToLLVM", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/TritonAMDGPUToLLVM/Passes.td", + deps = [":td_files"], +) + +gentbl_cc_library( + name = "triton_conversion_amdgpu_transforms_passes_inc_gen", + tbl_outs = {"include/TritonAMDGPUTransforms/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonAMDGPU", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/TritonAMDGPUTransforms/Passes.td", + deps = [":td_files"], +) + +cc_library( + name = "Analysis", + srcs = glob(["lib/Analysis/*.cpp"]), + hdrs = glob(["include/Analysis/*.h"]), + copts = _no_unused_variable + select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": ["-Wno-implicit-fallthrough"], + }), + includes = [ + "../../", + "include", + ], + deps = [ + ":TritonAMDGPU", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferIntRangeCommon", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:LoopLikeInterface", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:Support", + "@triton//:TritonDialects", + ], +) + +cc_library( + name = "TestAMDAnalysis", + srcs = glob(["test/lib/Analysis/*.cpp"]), + deps = [ + ":Analysis", + ":TritonAMDGPUToLLVM", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:TransformUtils", + "@triton//:TritonDialects", + "@triton//test:TritonTestAnalysis", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/amd/backend/BUILD b/third_party/amd/backend/BUILD new file mode 100644 --- /dev/null +++ b/third_party/amd/backend/BUILD @@ -0,0 +1,46 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") + +package( + # copybara:uncomment_begin + # default_applicable_licenses = ["@triton//:license"], + # default_compatible_with = [# "@build-target"], + # default_visibility = [ + # "//third_party/py/triton:__subpackages__", + # "@triton//:__subpackages__", + # ], + # copybara:uncomment_end_and_comment_begin + default_visibility = ["//visibility:public"], + # copybara:comment_end +) + +filegroup( + name = "files", + srcs = glob( + include = ["**/*.py"], + ), +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), +) + +cc_library( + name = "headers", + defines = [ + "__HIP_PLATFORM_AMD__", + ], + includes = [ + "include", + ], + textual_hdrs = glob([ + "include/**/*.h", + "include/**/*.hpp", + ]), +) + +exports_files([ + "include/TDMCommon.h", +]) diff --git b/third_party/f2reduce/BUILD b/third_party/f2reduce/BUILD new file mode 100644 --- /dev/null +++ b/third_party/f2reduce/BUILD @@ -0,0 +1,40 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") +# copybara:uncomment load("//tools/build_defs/license:license.bzl", "license") + +package( + # copybara:uncomment_begin + # default_applicable_licenses = ["@triton//:license"], + # default_compatible_with = [# "@build-target"], + # default_visibility = [ + # "@triton//:__subpackages__", + # ], + # copybara:uncomment_end_and_comment_begin + default_visibility = ["//visibility:public"], + # copybara:comment_end +) + +# copybara:uncomment_begin +# license( +# name = "license", +# license_text = "LICENCE.txt", +# ) +# +# licenses(["notice"]) +# +# exports_files(["LICENCE.txt"]) +# copybara:uncomment_end + +cc_library( + name = "f2reduce", + srcs = ["f2reduce.cpp"], + hdrs = ["f2reduce.h"], + # copybara:uncomment strip_include_prefix = "/third_party/triton", +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/nvidia/BUILD b/third_party/nvidia/BUILD new file mode 100644 --- /dev/null +++ b/third_party/nvidia/BUILD @@ -0,0 +1,565 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") + +package( + # copybara:uncomment_begin + # default_applicable_licenses = ["@triton//:license"], + # default_compatible_with = [# "@build-target"], + # default_visibility = [ + # "//third_party/py/enzyme_ad:__subpackages__", + # "//third_party/tensorflow/compiler/xla/backends/gpu:__subpackages__", + # "//third_party/tensorflow/compiler/xla/pjrt:__subpackages__", + # "@triton//:__subpackages__", + # ], + # copybara:uncomment_end_and_comment_begin + default_visibility = ["//visibility:public"], + # copybara:comment_end +) + +# copybara:uncomment_begin +# cc_library( +# name = "cublas_headers", +# hdrs = glob([ +# "include/*.h", +# ]), +# copts = ["-fexceptions"], +# features = ["-use_header_modules"], +# deps = ["@local_config_cuda//cuda:cuda_headers"], +# ) +# +# cc_library( +# name = "triton_nvidia", +# srcs = [ +# "triton_nvidia.cc", +# ], +# copts = [ +# "-fexceptions", +# "-Wno-sometimes-uninitialized", +# ], +# features = ["-use_header_modules"], +# visibility = [ +# "@triton//python:__subpackages__", +# ], +# deps = [ +# ":NVGPUDialect", +# ":NVGPUToLLVM", +# ":NVHopperTransforms", +# ":NVWSDialect", +# ":NVWSTransforms", +# ":TritonNVIDIAGPUToLLVM", +# ":cublas_headers", +# "@llvm-project//llvm:Core", +# "@llvm-project//llvm:Support", +# "@llvm-project//mlir:NVVMToLLVMIRTranslation", +# "@llvm-project//mlir:Pass", +# "@llvm-project//mlir:Transforms", +# "//third_party/nanobind", +# "//third_party/python_runtime:headers", +# "@triton//:TritonDialects", +# "@triton//:TritonGPUToLLVM", +# "@triton//:TritonNvidiaGPUTransforms", +# "@triton//python:passes", +# ], +# ) +# copybara:uncomment_end + +cc_library( + name = "NVGPUToLLVM", + srcs = glob([ + "lib/NVGPUToLLVM/*.cpp", + ]), + hdrs = glob([ + "include/NVGPUToLLVM/*.h", + ]), + copts = [ + "-Wno-unused-variable", + "-Wno-return-type", + ], + includes = [ + "..", + "include", + ], + deps = [ + ":NVGPUDialect", + ":TritonNVIDIAGPUToLLVM", + ":triton_conversion_nvgpu_to_llvm_passes_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:NVVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + ], +) + +cc_library( + name = "TritonNVIDIAGPUToLLVM", + srcs = glob( + include = [ + "lib/TritonNVIDIAGPUToLLVM/*.h", + "lib/TritonNVIDIAGPUToLLVM/**/*.cpp", + ], + exclude = ["lib/TritonNVIDIAGPUToLLVM/DecomposeUnsupportedConversions.cpp"], + ) + ["@triton//:instrument_to_llvm_srcs"], + hdrs = glob([ + "include/TritonNVIDIAGPUToLLVM/*.h", + "include/triton/Conversion/TritonGPUToLLVM/*.h", + ]) + [ + "lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAHelpers.h", + "lib/TritonNVIDIAGPUToLLVM/TargetInfo.h", + "lib/TritonNVIDIAGPUToLLVM/Utility.h", + ], + copts = [ + "-Wno-reorder-ctor", + "-Wno-unused-variable", + ], + includes = [ + "..", + "include", + "lib/TritonNVIDIAGPUToLLVM", + "lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM", + ], + deps = [ + ":NVGPUDialect", + ":triton_conversion_triton_nvidia_gpu_to_llvm_passes_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithToLLVM", + "@llvm-project//mlir:ArithTransforms", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:ControlFlowToLLVM", + "@llvm-project//mlir:ConversionPasses", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:GPUToNVVMTransforms", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:IndexDialect", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:MathToLLVM", + "@llvm-project//mlir:NVVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFToControlFlow", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBToLLVM", + "@triton//:GluonTransforms", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + "@triton//:TritonGPUTransforms", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//:TritonTools", + "@triton//:triton_gpu_attr_inc_gen", + ], +) + +gentbl_cc_library( + name = "triton_conversion_nvgpu_to_llvm_passes_inc_gen", + tbl_outs = {"include/NVGPUToLLVM/Passes.h.inc": [ + "--gen-pass-decls", + "--name=NVGPUToLLVM", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/NVGPUToLLVM/Passes.td", + deps = ["@triton//:td_files"], +) + +gentbl_cc_library( + name = "triton_conversion_triton_nvidia_gpu_to_llvm_passes_inc_gen", + tbl_outs = {"include/TritonNVIDIAGPUToLLVM/Passes.h.inc": [ + "--gen-pass-decls", + "--name=TritonNVIDIAGPUToLLVM", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/TritonNVIDIAGPUToLLVM/Passes.td", + deps = ["@triton//:td_files"], +) + +td_library( + name = "td_files", + srcs = glob([ + "include/Dialect/NVGPU/IR/*.td", + "include/Dialect/NVWS/IR/*.td", + ]), + includes = [ + "include", + ], + deps = [ + "@llvm-project//mlir:ArithOpsTdFiles", + "@llvm-project//mlir:CastInterfacesTdFiles", + "@llvm-project//mlir:ControlFlowInterfacesTdFiles", + "@llvm-project//mlir:DestinationStyleOpInterfaceTdFiles", + "@llvm-project//mlir:FunctionInterfacesTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", + "@llvm-project//mlir:LLVMOpsTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:PassBaseTdFiles", + "@llvm-project//mlir:SideEffectInterfacesTdFiles", + "@llvm-project//mlir:ViewLikeInterfaceTdFiles", + "@triton//:td_files", + ], +) + +gentbl_cc_library( + name = "nvgpu_ops_inc_gen", + tbl_outs = { + "include/Dialect/NVGPU/IR/OpsConversions.inc": ["--gen-llvmir-conversions"], + "include/Dialect/NVGPU/IR/Ops.h.inc": ["--gen-op-decls"], + "include/Dialect/NVGPU/IR/Ops.cpp.inc": ["--gen-op-defs"], + "include/Dialect/NVGPU/IR/OpsEnums.h.inc": ["--gen-enum-decls"], + "include/Dialect/NVGPU/IR/OpsEnums.cpp.inc": ["--gen-enum-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVGPU/IR/NVGPUOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvgpu_attr_inc_gen", + tbl_outs = { + "include/Dialect/NVGPU/IR/NVGPUAttrDefs.h.inc": ["--gen-attrdef-decls"], + "include/Dialect/NVGPU/IR/NVGPUAttrDefs.cpp.inc": ["--gen-attrdef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVGPU/IR/NVGPUAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvgpu_dialect_inc_gen", + tbl_outs = { + "include/Dialect/NVGPU/IR/Dialect.h.inc": ["--gen-dialect-decls"], + "include/Dialect/NVGPU/IR/Dialect.cpp.inc": ["--gen-dialect-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVGPU/IR/NVGPUDialect.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvws_dialect_inc_gen", + tbl_outs = [ + ( + ["--gen-dialect-decls"], + "include/Dialect/NVWS/IR/Dialect.h.inc", + ), + ( + ["--gen-dialect-defs"], + "include/Dialect/NVWS/IR/Dialect.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVWS/IR/NVWSDialect.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvws_attr_inc_gen", + tbl_outs = [ + ( + ["--gen-attrdef-decls"], + "include/Dialect/NVWS/IR/NVWSAttrDefs.h.inc", + ), + ( + ["--gen-attrdef-defs"], + "include/Dialect/NVWS/IR/NVWSAttrDefs.cpp.inc", + ), + ( + ["--gen-enum-decls"], + "include/Dialect/NVWS/IR/NVWSAttrEnums.h.inc", + ), + ( + ["--gen-enum-defs"], + "include/Dialect/NVWS/IR/NVWSAttrEnums.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVWS/IR/NVWSAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvws_types_inc_gen", + tbl_outs = [ + ( + ["--gen-typedef-decls"], + "include/Dialect/NVWS/IR/Types.h.inc", + ), + ( + ["--gen-typedef-defs"], + "include/Dialect/NVWS/IR/Types.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVWS/IR/NVWSTypes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvws_ops_inc_gen", + tbl_outs = [ + ( + ["--gen-op-decls"], + "include/Dialect/NVWS/IR/Ops.h.inc", + ), + ( + ["--gen-op-defs"], + "include/Dialect/NVWS/IR/Ops.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVWS/IR/NVWSOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvws_op_interfaces_inc_gen", + tbl_outs = [ + ( + ["--gen-op-interface-decls"], + "include/Dialect/NVWS/IR/NVWSOpInterfaces.h.inc", + ), + ( + ["--gen-op-interface-defs"], + "include/Dialect/NVWS/IR/NVWSOpInterfaces.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVWS/IR/NVWSOpInterfaces.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nvws_transforms_passes_inc_gen", + tbl_outs = [ + ( + [ + "--gen-pass-decls", + "--name=NVWSTransforms", + ], + "include/Dialect/NVWS/Transforms/Passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/Dialect/NVWS/Transforms/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "nv_hopper_transforms_passes_inc_gen", + tbl_outs = [ + ( + [ + "--gen-pass-decls", + "--name=NVHopperTransforms", + ], + "hopper/include/Transforms/Passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "hopper/include/Transforms/Passes.td", + deps = ["td_files"], +) + +cc_library( + name = "NVGPUDialect", + srcs = glob([ + "lib/Dialect/NVGPU/IR/*.cpp", + ]), + hdrs = glob([ + "include/Dialect/NVGPU/IR/*.h", + ]), + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + "-Wno-logical-op-parentheses", + ], + }), + includes = [ + "..", # because nvidia/include/Dialect/NVGPU/IR/Dialect.h.inc + "../..", # because third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h + "include", + ], + deps = [ + ":nvgpu_attr_inc_gen", + ":nvgpu_dialect_inc_gen", + ":nvgpu_ops_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InliningUtils", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:MathDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TensorDialect", + # The following is added to make Utility compile + "@triton//:TritonTools", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + ], +) + +cc_library( + name = "NVWSDialectHeader", + hdrs = ["include/Dialect/NVWS/IR/Dialect.h"], + includes = [ + "..", # because nvidia/include/Dialect/NVWS/IR/Dialect.h + ], + deps = [ + ":nvws_attr_inc_gen", + ":nvws_dialect_inc_gen", + ":nvws_op_interfaces_inc_gen", + ":nvws_ops_inc_gen", + ":nvws_types_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:TensorDialect", + "@triton//:TritonDialects", + ], +) + +cc_library( + name = "NVWSDialect", + srcs = glob([ + "lib/Dialect/NVWS/IR/*.cpp", + ]), + hdrs = glob([ + "include/Dialect/NVWS/IR/*.h", + ]), + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + includes = [ + "..", # because nvidia/include/Dialect/NVWS/IR/Dialect.h + "include", + ], + deps = [ + ":nvws_attr_inc_gen", + ":nvws_dialect_inc_gen", + ":nvws_op_interfaces_inc_gen", + ":nvws_ops_inc_gen", + ":nvws_types_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:TensorDialect", + "@triton//:TritonDialects", + "@triton//:TritonNvidiaGPUTransforms", + ], +) + +cc_library( + name = "NVWSTransforms", + srcs = glob([ + "lib/Dialect/NVWS/Transforms/*.cpp", + ]), + hdrs = glob([ + "include/Dialect/NVWS/Transforms/*.h", + "lib/Dialect/NVWS/Transforms/*.h", + ]), + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + "-Wno-string-conversion", + "-Wno-ctad-maybe-unsupported", + ], + }), + includes = [ + "../..", + "include", + ], + deps = [ + ":NVWSDialect", + ":nvws_transforms_passes_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferIntRangeCommon", + "@llvm-project//mlir:InferIntRangeInterface", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:NVVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TensorDialect", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@llvm-project//mlir:UBDialect", + "@triton//:TritonDialects", + "@triton//:TritonGPUTransforms", + ], +) + +cc_library( + name = "NVHopperTransforms", + srcs = glob([ + "hopper/lib/Transforms/**/*.cpp", + "hopper/lib/Transforms/**/*.h", + ]), + hdrs = glob([ + "hopper/include/Transforms/*.h", + ]), + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + "-Wno-string-conversion", + "-Wno-non-virtual-dtor", + "-Wno-reorder-ctor", + "-Wno-logical-op-parentheses", + ], + }), + includes = [ + "..", + "include", + ], + deps = [ + ":NVWSDialect", + ":nv_hopper_transforms_passes_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SCFToControlFlow", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@triton//:Dump", + "@triton//:GetEnv", + "@triton//:TritonDialects", + "@triton//:TritonGPUTransforms", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/nvidia/backend/BUILD b/third_party/nvidia/backend/BUILD new file mode 100644 --- /dev/null +++ b/third_party/nvidia/backend/BUILD @@ -0,0 +1,47 @@ +load("@pybind11_bazel//:build_defs.bzl", "pybind_extension") + +package( + default_applicable_licenses = ["@triton//:license"], + default_visibility = [ + "//third_party/py/triton:__subpackages__", + ], +) + +pybind_extension( + name = "cuda_utils", + srcs = ["driver.c"], + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + "-Wno-missing-braces", + "-Wno-sometimes-uninitialized", + ], + }), + visibility = [ + "@jaxlib/triton:__subpackages__", + "//third_party/py/triton:__subpackages__", + ], + deps = [ + # "@libcuda", + "//third_party/absl/cleanup", + "@local_config_cuda//cuda:cuda_headers", + "@local_config_cuda//cuda:cuda_runtime", + "@llvm-project//llvm:Support", + ], +) + +filegroup( + name = "files", + srcs = glob( + include = ["**/*.py"], + ), +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/nvidia/language/cuda/BUILD b/third_party/nvidia/language/cuda/BUILD new file mode 100644 --- /dev/null +++ b/third_party/nvidia/language/cuda/BUILD @@ -0,0 +1,21 @@ +package( + default_applicable_licenses = ["@triton//:license"], + default_visibility = [ + "//third_party/py/triton:__subpackages__", + ], +) + +filegroup( + name = "files", + srcs = glob( + include = ["**/*.py"], + ), +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/proton/BUILD b/third_party/proton/BUILD new file mode 100644 --- /dev/null +++ b/third_party/proton/BUILD @@ -0,0 +1,452 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") +load("@xla//third_party/rules_python/python:py_extension.bzl", "py_extension") +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") + +package( + # copybara:uncomment_begin + # default_applicable_licenses = ["@triton//:license"], + # default_compatible_with = [# "@build-target"], + # default_visibility = [ + # "@triton//:__subpackages__", + # ], + # copybara:uncomment_end_and_comment_begin + default_visibility = ["//visibility:public"], + # copybara:comment_end +) + +_no_unused_variable = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": ["-Wno-unused-variable"], +}) + +td_library( + name = "td_files", + srcs = glob(["Dialect/include/**/*.td"]), + includes = [ + "..", # because 'include proton/Dialect/include/' + "Dialect/include", + ], + deps = [ + "@llvm-project//mlir:OpBaseTdFiles", + "@triton//:td_files", + ], +) + +gentbl_cc_library( + name = "proton_ops_inc_gen", + tbl_outs = { + "Dialect/include/Dialect/Proton/IR/Dialect.h.inc": [ + "--gen-dialect-decls", + "--dialect=proton", + ], + "Dialect/include/Dialect/Proton/IR/Dialect.cpp.inc": [ + "--gen-dialect-defs", + "--dialect=proton", + ], + "Dialect/include/Dialect/Proton/IR/Ops.h.inc": ["--gen-op-decls"], + "Dialect/include/Dialect/Proton/IR/Ops.cpp.inc": ["--gen-op-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Dialect/Proton/IR/ProtonOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_attrdefs_inc_gen", + tbl_outs = { + "Dialect/include/Dialect/Proton/IR/OpsEnums.h.inc": ["--gen-enum-decls"], + "Dialect/include/Dialect/Proton/IR/OpsEnums.cpp.inc": ["--gen-enum-defs"], + "Dialect/include/Dialect/Proton/IR/ProtonAttrDefs.h.inc": ["--gen-attrdef-decls"], + "Dialect/include/Dialect/Proton/IR/ProtonAttrDefs.cpp.inc": ["--gen-attrdef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Dialect/Proton/IR/ProtonAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_gpu_ops_inc_gen", + tbl_outs = { + "Dialect/include/Dialect/ProtonGPU/IR/Dialect.h.inc": [ + "--gen-dialect-decls", + "--dialect=proton_gpu", + ], + "Dialect/include/Dialect/ProtonGPU/IR/Dialect.cpp.inc": [ + "--gen-dialect-defs", + "--dialect=proton_gpu", + ], + "Dialect/include/Dialect/ProtonGPU/IR/Ops.h.inc": ["--gen-op-decls"], + "Dialect/include/Dialect/ProtonGPU/IR/Ops.cpp.inc": ["--gen-op-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUOps.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_gpu_attrdefs_inc_gen", + tbl_outs = { + "Dialect/include/Dialect/ProtonGPU/IR/OpsEnums.h.inc": ["--gen-enum-decls"], + "Dialect/include/Dialect/ProtonGPU/IR/OpsEnums.cpp.inc": ["--gen-enum-defs"], + "Dialect/include/Dialect/ProtonGPU/IR/AttrDefs.h.inc": ["--gen-attrdef-decls"], + "Dialect/include/Dialect/ProtonGPU/IR/AttrDefs.cpp.inc": ["--gen-attrdef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUAttrDefs.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_gpu_types_inc_gen", + tbl_outs = { + "Dialect/include/Dialect/ProtonGPU/IR/Types.h.inc": ["--gen-typedef-decls"], + "Dialect/include/Dialect/ProtonGPU/IR/Types.cpp.inc": ["--gen-typedef-defs"], + }, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUTypes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_to_proton_gpu_inc_gen", + tbl_outs = {"Dialect/include/Conversion/ProtonToProtonGPU/Passes.h.inc": [ + "--gen-pass-decls", + "--name=ProtonToProtonGPU", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Conversion/ProtonToProtonGPU/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_gpu_transforms_inc_gen", + tbl_outs = {"Dialect/include/Dialect/ProtonGPU/Transforms/Passes.h.inc": [ + "--gen-pass-decls", + "--name=ProtonGPU", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Dialect/ProtonGPU/Transforms/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_gpu_to_llvm_inc_gen", + tbl_outs = {"Dialect/include/Conversion/ProtonGPUToLLVM/Passes.h.inc": [ + "--gen-pass-decls", + "--name=ProtonGPUToLLVM", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Conversion/ProtonGPUToLLVM/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_amd_gpu_to_llvm_inc_gen", + tbl_outs = {"Dialect/include/Conversion/ProtonGPUToLLVM/ProtonAMDGPUToLLVM/Passes.h.inc": [ + "--gen-pass-decls", + "--name=ProtonAMDGPUToLLVM", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Conversion/ProtonGPUToLLVM/ProtonAMDGPUToLLVM/Passes.td", + deps = ["td_files"], +) + +gentbl_cc_library( + name = "proton_nvidia_gpu_to_llvm_inc_gen", + tbl_outs = {"Dialect/include/Conversion/ProtonGPUToLLVM/ProtonNvidiaGPUToLLVM/Passes.h.inc": [ + "--gen-pass-decls", + "--name=ProtonNvidiaGPUToLLVM", + ]}, + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/include/Conversion/ProtonGPUToLLVM/ProtonNvidiaGPUToLLVM/Passes.td", + deps = ["td_files"], +) + +cc_library( + name = "ProtonIR", + srcs = glob([ + "Dialect/lib/Dialect/Proton/IR/*.cpp", + ]), + hdrs = glob([ + "Dialect/include/Dialect/Proton/IR/*.h", + ]), + includes = [ + "..", # because '#include "proton/Dialect/include/"' + "Dialect/include", + ], + deps = [ + ":proton_attrdefs_inc_gen", + ":proton_ops_inc_gen", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InliningUtils", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Support", + "@triton//:TritonDialects", + ], +) + +cc_library( + name = "ProtonGPUIR", + srcs = glob([ + "Dialect/lib/Dialect/ProtonGPU/IR/*.cpp", + ]), + hdrs = glob([ + "Dialect/include/Dialect/ProtonGPU/IR/*.h", + ]), + copts = _no_unused_variable, + includes = [ + "..", # because '#include "proton/Dialect/include/"' + "Dialect/include", + ], + deps = [ + ":ProtonIR", + ":proton_gpu_attrdefs_inc_gen", + ":proton_gpu_ops_inc_gen", + ":proton_gpu_types_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InliningUtils", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Support", + "@triton//:TritonDialects", + ], +) + +cc_library( + name = "ProtonAnalysis", + srcs = glob(["Dialect/lib/Analysis/*.cpp"]), + hdrs = glob(["Dialect/include/Analysis/*.h"]), + includes = ["Dialect/include"], + deps = [ + ":ProtonIR", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", + "@triton//:TritonDialects", + ], +) + +cc_library( + name = "ProtonToProtonGPU", + srcs = glob(["Dialect/lib/ProtonToProtonGPU/*.cpp"]), + hdrs = glob(["Dialect/include/Conversion/ProtonToProtonGPU/*.h"]), + copts = _no_unused_variable, + deps = [ + ":ProtonAnalysis", + ":ProtonGPUIR", + ":ProtonIR", + ":proton_to_proton_gpu_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + ], +) + +cc_library( + name = "ProtonGPUTransforms", + srcs = glob(["Dialect/lib/Dialect/ProtonGPU/Transforms/*.cpp"]), + hdrs = glob(["Dialect/include/Dialect/ProtonGPU/Transforms/*.h"]), + includes = ["Dialect/include"], + deps = [ + ":ProtonGPUIR", + ":proton_gpu_transforms_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InliningUtils", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:Pass", + ], +) + +cc_library( + name = "ProtonGPUToLLVM", + srcs = glob([ + "Dialect/lib/ProtonGPUToLLVM/*.cpp", + "Dialect/lib/ProtonGPUToLLVM/**/*.cpp", + ]), + hdrs = glob([ + "Dialect/include/Conversion/ProtonGPUToLLVM/*.h", + "Dialect/include/Conversion/ProtonGPUToLLVM/**/*.h", + ]), + copts = _no_unused_variable, + deps = [ + ":ProtonGPUIR", + ":proton_amd_gpu_to_llvm_inc_gen", + ":proton_gpu_to_llvm_inc_gen", + ":proton_nvidia_gpu_to_llvm_inc_gen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:AMDGPUUtils", + "@llvm-project//mlir:ArithToLLVM", + "@llvm-project//mlir:ControlFlowDialect", + "@llvm-project//mlir:ControlFlowToLLVM", + "@llvm-project//mlir:ConversionPasses", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:GPUToNVVMTransforms", + "@llvm-project//mlir:GPUToROCDLTransforms", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMCommonConversion", + "@llvm-project//mlir:LLVMDialect", + "@llvm-project//mlir:NVVMDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:ROCDLDialect", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@triton//:TritonDialects", + "@triton//:TritonGPUToLLVM", + "@triton//third_party/amd:TritonAMDGPU", + "@triton//third_party/amd:TritonAMDGPUToLLVM", + "@triton//third_party/nvidia:TritonNVIDIAGPUToLLVM", + ], +) + +genrule( + name = "fixed_triton_proton_cc", + srcs = ["Dialect/triton_proton.cc"], + outs = ["Dialect/triton_proton_fixed.cc"], + cmd = "sed -e 's/PYBIND11_MODULE/NB_MODULE/' -e 's/init_triton_proton(std::move(m))/init_triton_proton(m)/' $< > $@", +) + +py_extension( + name = "triton_proton", + srcs = [ + ":fixed_triton_proton_cc", + ], + copts = ["-fexceptions"], + features = ["-use_header_modules"], + visibility = [ + "//third_party/py/triton:__subpackages__", + ], + deps = [ + ":ProtonAnalysis", + ":ProtonGPUIR", + ":ProtonGPUToLLVM", + ":ProtonGPUTransforms", + ":ProtonIR", + ":ProtonToProtonGPU", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "//third_party/nanobind", + "//third_party/python_runtime:headers", + "@triton//python:ir", + "@triton//python:passes", + ], +) + +cc_library( + name = "instrumentation_metadata_hdr", + includes = ["csrc/include"], + textual_hdrs = ["csrc/include/Profiler/Instrumentation/InstrumentationMetadata.h"], +) + +genrule( + name = "device_type_hdr_gen", + srcs = ["common/include/DeviceType.in"], + outs = ["common/include/DeviceType.h"], + cmd = "sed 's/@BACKEND_DEVICE_TAGS@//g' $< > $@", +) + +genrule( + name = "registered_backends_src_gen", + srcs = ["csrc/lib/Backend/RegisteredBackends.in"], + outs = ["csrc/lib/Backend/RegisteredBackends.cpp"], + cmd = "sed -e 's/@BACKEND_REGISTRATION_HOOKS_FORWARD_DECLS@//g' -e 's/@BACKEND_REGISTRATION_HOOKS@//g' $< > $@", +) + +cc_library( + name = "proton_cxx", + srcs = glob( + [ + "csrc/**/*.cpp", + "common/**/*.cpp", + ], + exclude = [ + "csrc/lib/Data/Dump/Hatchet.cpp", + "csrc/lib/Data/Dump/HatchetMsgPack.cpp", + "csrc/lib/Data/Dump/TreeDataDump.cpp", + ], + ) + [ + ":device_type_hdr_gen", + ":registered_backends_src_gen", + ], + hdrs = glob( + [ + "csrc/**/*.h", + "common/**/*.h", + ], + exclude = ["csrc/include/Profiler/Instrumentation/InstrumentationMetadata.h"], + ), + copts = [ + "-Wno-reorder-ctor", + "-Wno-implicit-fallthrough", + "-Wno-pessimizing-move", + "-Wno-unused-variable", + "-fexceptions", + ], + features = [ + "exceptions", + "-use_header_modules", + ], + includes = [ + "common/include", + "csrc/include", + "csrc/lib/Data", + ], + textual_hdrs = [ + "csrc/lib/Data/Dump/Hatchet.cpp", + "csrc/lib/Data/Dump/HatchetMsgPack.cpp", + "csrc/lib/Data/Dump/TreeDataDump.cpp", + ], + deps = [ + ":instrumentation_metadata_hdr", + "@local_config_cuda//cuda:cuda_headers", + "//third_party/json", + "//third_party/nanobind", + "//third_party/python_runtime:headers", + "@triton//third_party/amd/backend:headers", + ], +) + +py_extension( + name = "libproton", + srcs = ["csrc/Proton.cpp"], + copts = [ + "-fexceptions", + "-Wno-pessimizing-move", + ], + features = ["-use_header_modules"], + visibility = [ + "//third_party/py/triton:__subpackages__", + ], + deps = [ + ":proton_cxx", + "//third_party/nanobind", + "//third_party/python_runtime:headers", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/proton/proton/BUILD b/third_party/proton/proton/BUILD new file mode 100644 --- /dev/null +++ b/third_party/proton/proton/BUILD @@ -0,0 +1,21 @@ +package( + default_applicable_licenses = ["@triton//:license"], + default_visibility = [ + "//third_party/py/triton:__subpackages__", + ], +) + +filegroup( + name = "files", + srcs = glob( + include = ["**/*.py"], + ), +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/third_party/proton/test/BUILD b/third_party/proton/test/BUILD new file mode 100644 --- /dev/null +++ b/third_party/proton/test/BUILD @@ -0,0 +1,40 @@ +load("//third_party/py/pytest:pytest_defs.bzl", "pytest_multi_tests") + +package( + default_applicable_licenses = ["@triton//:license"], +) + +pytest_multi_tests( + name = "all_tests", + tags = ["requires-gpu-sm90"], + target_compatible_with = select( + {"@rules_cuda//cuda:is_enabled": []}, + no_match_error = "Requires --config=cuda", + ), + tests = glob( + include = ["test_*.py"], + exclude = [ + "test_viewer.py", # Requires hatchet as dependency, which we do not have imported. + "test_cmd.py", # Tries to do filepath magic, would need patching to work internally. + "test_override.py", # Tries to do filepath magic, would need patching to work internally. + "test_profile.py", # Needs to dlopen cupti, but we only have a static version of it. + "test_api.py", # Needs to dlopen cupti, but we only have a static version of it. + "test_lib.py", # Needs to dlopen cupti, but we only have a static version of it. + "test_instrumentation.py", # Needs to dlopen cupti, but we only have a static version of it. + ], + ), + deps = [ + "//third_party/py/pandas", + "//third_party/py/torch:pytorch", + "//third_party/py/triton", + "//third_party/py/triton/profiler", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +) diff --git b/unittest/BUILD b/unittest/BUILD new file mode 100644 --- /dev/null +++ b/unittest/BUILD @@ -0,0 +1,158 @@ +load("@rules_cc//cc:cc_test.bzl", "cc_test") +load("//tools/build_defs/build_test:build_test.bzl", "build_test") + +package( + default_applicable_licenses = ["@triton//:license"], + default_compatible_with = [# "@build-target"], + default_visibility = ["@triton//:__subpackages__"], +) + +cc_test( + name = "AnalysisTest", + srcs = glob(["Analysis/*.cpp"]), + deps = [ + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + "@triton//:TritonDialects", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//:TritonTransforms", + ], +) + +cc_test( + name = "DialectTestCatchAll", + srcs = glob( + [ + "Dialect/**/*.cpp", + ], + exclude = [ + "Dialect/TritonGPU/DialectTest.cpp", + "Dialect/TritonGPU/LinearLayoutConversionsTest.cpp", + "Dialect/TritonGPU/SwizzleTest.cpp", + ], + ), + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + deps = [ + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:AsmParser", + "@llvm-project//mlir:IR", + "@triton//:TritonDialects", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//:TritonTools", + "@triton//:TritonTransforms", + ], +) + +cc_test( + name = "DialectTest", + srcs = [ + "Dialect/TritonGPU/DialectTest.cpp", + ], + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + deps = [ + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:AsmParser", + "@llvm-project//mlir:IR", + "@triton//:TritonDialects", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//:TritonTools", + "@triton//:TritonTransforms", + ], +) + +cc_test( + name = "LinearLayoutConversionsTest", + srcs = [ + "Dialect/TritonGPU/LinearLayoutConversionsTest.cpp", + ], + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + deps = [ + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:AsmParser", + "@llvm-project//mlir:IR", + "@triton//:TritonDialects", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//:TritonTools", + "@triton//:TritonTransforms", + ], +) + +cc_test( + name = "SwizzleTest", + srcs = [ + "Dialect/TritonGPU/SwizzleTest.cpp", + ], + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + deps = [ + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:AsmParser", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", + "@triton//:TritonDialects", + "@triton//:TritonNvidiaGPUTransforms", + "@triton//:TritonTools", + "@triton//:TritonTransforms", + ], +) + +cc_test( + name = "ToolsTest", + srcs = glob(["Tools/*.cpp"]), + copts = select({ + "@triton//:compiler_is_msvc": [], + "//conditions:default": [ + "-Wno-unused-variable", + ], + }), + deps = [ + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Support", + "@triton//:TritonTools", + ], +) + +build_test( + name = "build_test", + allow_empty_target = False, + targets = [ + ":ToolsTest", + ":SwizzleTest", + ":LinearLayoutConversionsTest", + ":DialectTestCatchAll", + ":AnalysisTest", + ":DialectTest", + ], +) + +filegroup( + name = "all-files", + srcs = glob( + include = ["**"], + ), + visibility = ["@triton//:__subpackages__"], +)