diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 0acdb3133c76..1ab571315df5 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -8120,7 +8120,7 @@ device_check: NoCheck device_guard: False dispatch: - ZeroTensor, Meta, CPU, CUDA, QuantizedCPU, QuantizedCUDA, MPS: view + ZeroTensor, Meta, CPU, CUDA, QuantizedCPU, QuantizedCUDA, MPS, MTIA: view MkldnnCPU: mkldnn_view NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: view_nested tags: core diff --git a/build.bzl b/build.bzl index f8bc5acefa6f..28a7cedbee3b 100644 --- a/build.bzl +++ b/build.bzl @@ -79,6 +79,8 @@ def define_targets(rules): aten_ufunc_generated_cuda_sources() ) + gen_aten_outs_mtia = GENERATED_H_MTIA + GENERATED_CPP_MTIA + gen_aten_outs = ( GENERATED_H + GENERATED_H_CORE + GENERATED_CPP + GENERATED_CPP_CORE + @@ -86,7 +88,7 @@ def define_targets(rules): aten_ufunc_generated_cpu_sources() + aten_ufunc_generated_cpu_kernel_sources() + [ "Declarations.yaml", - ] + gen_aten_outs_cuda + ] + gen_aten_outs_cuda + gen_aten_outs_mtia ) rules.genrule( @@ -208,6 +210,15 @@ GENERATED_CPP_CUDA = [ "RegisterQuantizedCUDA_0.cpp", ] +GENERATED_H_MTIA = [ + "MTIAFunctions.h", + "MTIAFunctions_inl.h", +] + +GENERATED_CPP_MTIA = [ + "RegisterMTIA_0.cpp", +] + GENERATED_CPP = [ "Functions.cpp", "RegisterBackendSelect.cpp", diff --git a/torchgen/gen.py b/torchgen/gen.py index 1c27d32be727..94ccf62ccb37 100644 --- a/torchgen/gen.py +++ b/torchgen/gen.py @@ -2977,6 +2977,7 @@ def main() -> None: DispatchKey.CompositeExplicitAutograd, DispatchKey.CompositeExplicitAutogradNonFunctional, DispatchKey.Meta, + DispatchKey.MTIA, } aoti_backends = { diff --git a/torchgen/model.py b/torchgen/model.py index 4715759d322e..a3c0542788bb 100644 --- a/torchgen/model.py +++ b/torchgen/model.py @@ -305,6 +305,7 @@ dispatch_keys = [ DispatchKey.QuantizedMeta, DispatchKey.NestedTensorMeta, DispatchKey.ZeroTensor, + DispatchKey.MTIA, ]