feat(musa): 13 reduction/normalization mudnn kernels + unbreak MUSA build

csrc/aten/backends/musa/generated/musa_register.inc

@@ -9,14 +9,21 @@
 // check, whereas an unregistered op simply reaches the cpu_fallback. So this
 // file is exactly the MUSA coverage set.
 
+  m.impl("_log_softmax", WrapperPrivLogSoftmax);
+  m.impl("_log_softmax_backward_data", WrapperPrivLogSoftmaxBackwardData);
   m.impl("_softmax", WrapperPrivSoftmax);
+  m.impl("_softmax_backward_data", WrapperPrivSoftmaxBackwardData);
   m.impl("abs", WrapperAbs);
   m.impl("acos", WrapperAcos);
   m.impl("add.Scalar", WrapperAddScalar);
   m.impl("add.Tensor", WrapperAddTensor);
   m.impl("addcdiv", WrapperAddcdiv);
   m.impl("addcmul", WrapperAddcmul);
   m.impl("addmm", WrapperAddmm);
+  m.impl("all.dim", WrapperAllDim);
+  m.impl("amax", WrapperAmax);
+  m.impl("amin", WrapperAmin);
+  m.impl("any.dim", WrapperAnyDim);
   m.impl("atan", WrapperAtan);
   m.impl("baddbmm", WrapperBaddbmm);
   m.impl("bmm", WrapperBmm);
@@ -54,6 +61,7 @@
   m.impl("le.Tensor", WrapperLeTensor);
   m.impl("leaky_relu", WrapperLeakyRelu);
   m.impl("leaky_relu_backward", WrapperLeakyReluBackward);
+  m.impl("linalg_vector_norm", WrapperLinalgVectorNorm);
   m.impl("log", WrapperLog);
   m.impl("log10", WrapperLog10);
   m.impl("log1p", WrapperLog1p);
@@ -72,11 +80,14 @@
   m.impl("mm.out", WrapperMmOut);
   m.impl("mul.Scalar", WrapperMulScalar);
   m.impl("mul.Tensor", WrapperMulTensor);
+  m.impl("native_layer_norm", WrapperNativeLayerNorm);
+  m.impl("native_layer_norm_backward", WrapperNativeLayerNormBackward);
   m.impl("ne.Scalar", WrapperNeScalar);
   m.impl("ne.Tensor", WrapperNeTensor);
   m.impl("neg", WrapperNeg);
   m.impl("pow.Tensor_Scalar", WrapperPowTensorScalar);
   m.impl("pow.Tensor_Tensor", WrapperPowTensorTensor);
+  m.impl("prod.dim_int", WrapperProdDimInt);
   m.impl("reciprocal", WrapperReciprocal);
   m.impl("relu", WrapperRelu);
   m.impl("remainder.Scalar", WrapperRemainderScalar);
@@ -91,6 +102,7 @@
   m.impl("sin", WrapperSin);
   m.impl("softplus", WrapperSoftplus);
   m.impl("sqrt", WrapperSqrt);
+  m.impl("std.correction", WrapperStdCorrection);
   m.impl("sub.Scalar", WrapperSubScalar);
   m.impl("sub.Tensor", WrapperSubTensor);
   m.impl("sum", WrapperSum);
@@ -100,4 +112,5 @@
   m.impl("tanh_backward", WrapperTanhBackward);
   m.impl("threshold_backward", WrapperThresholdBackward);
   m.impl("trunc", WrapperTrunc);
+  m.impl("var.correction", WrapperVarCorrection);
   m.impl("where.self", WrapperWhereSelf);

csrc/aten/backends/musa/mudnn_common.h

@@ -222,20 +222,28 @@ inline mudnn::MemoryMaintainer MudnnWorkspaceFor(const at::Tensor& reference) {
   };
 }
 
-// True when mudnn's Reduce would fault on this input rather than return a
-// status. Measured on mudnn v3300: a Reduce over *more than one* dim whose input
-// is fully 0-strided (every dim with extent > 1 has stride 0, i.e. the whole
-// tensor is a broadcast of a single element) raises SIGFPE inside the vendor
-// library -- an uncatchable crash, not a NOT_SUPPORTED status. A single-dim
-// reduce over the same input is fine, and so is a multi-dim reduce as soon as
-// any one stride is non-zero.
+// True when mudnn's Reduce misbehaves on this input and it must be materialized
+// first. The trigger is an input that is a *broadcast of a single element* --
+// every dim with extent > 1 has stride 0 -- which reaches real code through bias
+// gradients: `linear(x, w, b).sum()` backward reduces a grad_output that
+// autograd produced as `ones.expand(...)`, whose storage is one float.
 //
-// Kernels check this and materialize the input with a contiguous copy first. It
-// reaches real code through convolution bias gradients: `conv(x, w, b).sum()`
-// backward reduces a grad_output that autograd produced with `ones.expand()`.
-inline bool MudnnReduceWouldFault(
-    const at::Tensor& self, size_t num_reduced_dims) {
-  if (num_reduced_dims <= 1) {
+// Two distinct failures were measured on mudnn v3300, both silent in the status:
+//
+//   - Reducing over *more than one* dim raises SIGFPE inside the vendor library
+//     -- an uncatchable crash, not a NOT_SUPPORTED status.
+//   - Reducing over a *single* dim intermittently writes only out[0] and leaves
+//     the remaining output elements untouched, so the answer is whatever the
+//     caching allocator last left in that block. Observed as a wrong bias
+//     gradient (`[4, 34, 38, 42, 46]` where elements 1.. were a previous op's
+//     result); the same reduce on a materialized copy is always correct.
+//
+// A multi-dim reduce is fine as soon as any one stride is non-zero, but the
+// single-dim partial write does not reproduce standalone, so this predicate does
+// not try to be narrower than "fully 0-strided". The copy it forces is one
+// element wide on input, so the cost is a single small materialization.
+inline bool MudnnReduceNeedsContiguous(const at::Tensor& self) {
+  if (self.is_contiguous()) {
     return false;
   }
   for (int64_t d = 0; d < self.dim(); ++d) {

csrc/runtime/guard.h

@@ -18,9 +18,11 @@
 // exports c10::hip::HIPCachingAllocator with zero c10::cuda symbols (see
 // backends/dcu_memory.h), so c10::cuda::getCurrentCUDAStream would fail to
 // resolve there even though DCU's CUDA-compat runtime satisfies plain
-// cudaStream_t calls.
+// cudaStream_t calls. MUSA is excluded for the plainer reason that the Moore
+// Threads toolkit ships no CUDA runtime at all, so the header itself is absent
+// -- same exclusion as hooks.h and copy_ops.cc already carry.
 #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_DCU) && \
-    !defined(USE_GCU)
+    !defined(USE_GCU) && !defined(USE_MUSA)
 #define FLAGOS_GUARD_HAS_CUDA_STREAM 1
 #include <c10/cuda/CUDAStream.h>
 #else