diff --git a/musa_ext/graph/fusion/tensordot_fusion.cc b/musa_ext/graph/fusion/tensordot_fusion.cc index 34a1e795..11df4e6a 100644 --- a/musa_ext/graph/fusion/tensordot_fusion.cc +++ b/musa_ext/graph/fusion/tensordot_fusion.cc @@ -821,6 +821,40 @@ Status MusaTensorDotFusion::Apply(GraphDef* graph, fuse_node_names.erase(name); } + // 递归保留 shared node 的上游依赖；否则虽然 shared node 本身会保留， + // 但它依赖的内部 shape/perm 辅助节点仍可能被删掉，留下悬空输入。 + std::vector keep_worklist(shared_nodes.begin(), + shared_nodes.end()); + std::unordered_set preserved_dependencies(shared_nodes.begin(), + shared_nodes.end()); + while (!keep_worklist.empty()) { + const std::string node_name = keep_worklist.back(); + keep_worklist.pop_back(); + + const int node_idx = FusionGraphUtils::FindNodeIndex(*graph, node_name); + if (node_idx < 0 || node_idx >= graph->node_size()) { + continue; + } + + const NodeDef& node = graph->node(node_idx); + for (int input_idx = 0; input_idx < node.input_size(); ++input_idx) { + const std::string producer = + FusionGraphUtils::GetProducerNodeName(node.input(input_idx)); + if (producer.empty() || producer == output_name || + producer == input_a_name || producer == input_b_name || + !fuse_node_names.count(producer) || + preserved_dependencies.count(producer)) { + continue; + } + + VLOG(2) << "[TensorDot::Apply] keeping dependency of shared node: " + << producer << " -> " << node_name; + fuse_node_names.erase(producer); + preserved_dependencies.insert(producer); + keep_worklist.push_back(producer); + } + } + VLOG(2) << "[TensorDot::Apply] will remove " << fuse_node_names.size() << " fused sub-graph nodes"; diff --git a/musa_ext/kernels/array/musa_matrix_set_diag_kernel.mu b/musa_ext/kernels/array/musa_matrix_set_diag_kernel.mu index ef71c969..d8cf3107 100644 --- a/musa_ext/kernels/array/musa_matrix_set_diag_kernel.mu +++ b/musa_ext/kernels/array/musa_matrix_set_diag_kernel.mu @@ -31,7 +31,7 @@ __global__ void MatrixSetDiagKernel( const int batch = batch_and_diag_index / num_diags; const int diag_index_in_input = batch_and_diag_index - batch * num_diags; const int diag_index = upper_diag_index - diag_index_in_input; - index_in_the_diagonal = + index_in_the_diagonal += ComputeContentOffset(diag_index, max_diag_len, m, n, left_align_superdiagonal, left_align_subdiagonal); const int y_index = index_in_the_diagonal - min(0, diag_index); @@ -112,7 +112,7 @@ void MusaMatrixSetDiagKernelLauncher( musaStream_t, typename TTypes::ConstTensor&, typename TTypes::ConstTensor&, \ typename TTypes::Tensor&, const Eigen::Index, const Eigen::Index, const Eigen::Index, \ const bool, const bool); - + INSTANTIATE_MATRIX_SET_DIAG_KERNEL(float); INSTANTIATE_MATRIX_SET_DIAG_KERNEL(double); INSTANTIATE_MATRIX_SET_DIAG_KERNEL(int32_t); @@ -122,4 +122,4 @@ INSTANTIATE_MATRIX_SET_DIAG_KERNEL(Eigen::bfloat16); INSTANTIATE_MATRIX_SET_DIAG_KERNEL(bool); } // namespace musa -} // namespace tensorflow \ No newline at end of file +} // namespace tensorflow diff --git a/musa_ext/kernels/array/musa_matrix_set_diag_op.cc b/musa_ext/kernels/array/musa_matrix_set_diag_op.cc index ca91ecee..ef6669ee 100644 --- a/musa_ext/kernels/array/musa_matrix_set_diag_op.cc +++ b/musa_ext/kernels/array/musa_matrix_set_diag_op.cc @@ -20,7 +20,7 @@ void ReadAlignment(OpKernelConstruction* context, OP_REQUIRES_OK(context, context->GetAttr("align", &align)); *left_align_superdiagonal = align == "LEFT_LEFT" || align == "LEFT_RIGHT"; - *left_align_subdiagonal = align == "RIGHT_LEFT" || align == "RIGHT_RIGHT"; + *left_align_subdiagonal = align == "LEFT_LEFT" || align == "RIGHT_LEFT"; } template @@ -159,7 +159,7 @@ class MusaMatrixSetDiagOp : public MusaOpKernel { musaStream_t stream = GetMusaStreamByCtx(context); MusaMatrixSetDiagKernelLauncher( stream, input_reshaped, diag_reshaped, output_reshaped, - lower_diag_index, upper_diag_index, num_diags, + lower_diag_index, upper_diag_index, max_diag_len, left_align_superdiagonal_, left_align_subdiagonal_); } @@ -171,15 +171,15 @@ class MusaMatrixSetDiagOp : public MusaOpKernel { static constexpr int kNumV1Inputs = 2; }; -#define REGISTER_MATRIX_SET_DIAG_KERNEL(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("MusaMatrixSetDiag").Device("MUSA").TypeConstraint("T"), \ - MusaMatrixSetDiagOp); \ - REGISTER_KERNEL_BUILDER( \ - Name("MusaMatrixSetDiagV2").Device("MUSA").TypeConstraint("T"), \ - MusaMatrixSetDiagOp); \ - REGISTER_KERNEL_BUILDER( \ - Name("MusaMatrixSetDiagV3").Device("MUSA").TypeConstraint("T"), \ +#define REGISTER_MATRIX_SET_DIAG_KERNEL(type) \ + REGISTER_KERNEL_BUILDER( \ + Name("MatrixSetDiag").Device("MUSA").TypeConstraint("T"), \ + MusaMatrixSetDiagOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("MatrixSetDiagV2").Device("MUSA").TypeConstraint("T"), \ + MusaMatrixSetDiagOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("MatrixSetDiagV3").Device("MUSA").TypeConstraint("T"), \ MusaMatrixSetDiagOp); REGISTER_MATRIX_SET_DIAG_KERNEL(float); @@ -192,4 +192,4 @@ REGISTER_MATRIX_SET_DIAG_KERNEL(bool); #undef REGISTER_MATRIX_SET_DIAG_KERNEL } // namespace musa -} // namespace tensorflow \ No newline at end of file +} // namespace tensorflow diff --git a/musa_ext/kernels/control_flow/musa_switch_op.cc b/musa_ext/kernels/control_flow/musa_switch_op.cc index afdfb5c6..39eb9a34 100644 --- a/musa_ext/kernels/control_flow/musa_switch_op.cc +++ b/musa_ext/kernels/control_flow/musa_switch_op.cc @@ -28,9 +28,9 @@ class MusaSwitchOp : public MusaOpKernel { } }; // class MusaSwitchOp -#define REGISTER_MUSA_SWITCH(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("MusaSwitch").Device(DEVICE_MTGPU).TypeConstraint("T"), \ +#define REGISTER_MUSA_SWITCH(type) \ + REGISTER_KERNEL_BUILDER( \ + Name("Switch").Device(DEVICE_MTGPU).TypeConstraint("T"), \ MusaSwitchOp); REGISTER_MUSA_SWITCH(float); @@ -43,4 +43,4 @@ REGISTER_MUSA_SWITCH(uint8); REGISTER_MUSA_SWITCH(bool); } // namespace musa -} // namespace tensorflow \ No newline at end of file +} // namespace tensorflow diff --git a/musa_ext/kernels/fusion/musa_sigmoid_calibration_op.cc b/musa_ext/kernels/fusion/musa_sigmoid_calibration_op.cc index f4a54bfd..7180b125 100644 --- a/musa_ext/kernels/fusion/musa_sigmoid_calibration_op.cc +++ b/musa_ext/kernels/fusion/musa_sigmoid_calibration_op.cc @@ -8,6 +8,13 @@ namespace tensorflow { namespace musa { +REGISTER_OP("MusaSigmoidCalibration") + .Input("input: T") + .Input("scale: T") + .Output("output: T") + .Attr("T: {float, double, half, bfloat16}") + .SetShapeFn(shape_inference::UnchangedShape); + template void LaunchSigmoidCalibrationKernel(const void*, const void*, void*, int, musaStream_t); diff --git a/test/ops/switch_op_test.py b/test/ops/switch_op_test.py index fad3532d..9eb254f5 100755 --- a/test/ops/switch_op_test.py +++ b/test/ops/switch_op_test.py @@ -24,9 +24,8 @@ class SwitchOpTest(MUSATestCase): """Tests for tf.raw_ops.Switch on MUSA.""" - def _run_on_device(self, data, pred, device): - with tf.device(device): - return tf.raw_ops.Switch(data=data, pred=pred) + def _run_branch(self, data, pred, branch_index): + return tf.raw_ops.Switch(data=data, pred=pred)[branch_index] def _compare_active_branch(self, dtype, pred_value): np_dtype = dtype.as_numpy_dtype @@ -40,18 +39,32 @@ def _compare_active_branch(self, dtype, pred_value): data = tf.constant(data_np, dtype=dtype) pred = tf.constant(pred_value) - cpu_false, cpu_true = self._run_on_device(data, pred, '/CPU:0') - musa_false, musa_true = self._run_on_device(data, pred, '/device:MUSA:0') - rtol = 1e-2 if dtype in [tf.float16, tf.bfloat16] else 1e-5 atol = 1e-2 if dtype in [tf.float16, tf.bfloat16] else 1e-8 - if pred_value: - # Only true branch should carry the payload. - self.assertAllClose(cpu_true.numpy(), musa_true.numpy(), rtol=rtol, atol=atol) - else: - # Only false branch should carry the payload. - self.assertAllClose(cpu_false.numpy(), musa_false.numpy(), rtol=rtol, atol=atol) + active_branch = 1 if pred_value else 0 + self._compare_cpu_musa_results( + lambda input_data, input_pred: self._run_branch( + input_data, input_pred, active_branch), + [data, pred], + dtype, + rtol=rtol, + atol=atol) + self._compare_cpu_musa_results( + lambda input_data, input_pred: self._run_branch( + input_data, input_pred, 1 - active_branch), + [data, pred], + dtype, + rtol=rtol, + atol=atol) + + def _assert_invalid_pred_on_device(self, data, pred, device): + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, "must be a scalar"): + self._test_op_device_placement( + lambda input_data, input_pred: tf.raw_ops.Switch( + data=input_data, pred=input_pred), + [data, pred], + device) def test_switch_true_branch(self): for dtype in [tf.float32, tf.float16, tf.int32, tf.bool]: @@ -62,10 +75,11 @@ def test_switch_false_branch(self): self._compare_active_branch(dtype, False) def test_pred_must_be_scalar(self): - with self.assertRaisesRegex(tf.errors.InvalidArgumentError, "must be a scalar"): - with tf.device('/device:MUSA:0'): - tf.raw_ops.Switch(data=tf.constant([1, 2, 3], dtype=tf.int32), - pred=tf.constant([True, False], dtype=tf.bool)) + data = tf.constant([1, 2, 3], dtype=tf.int32) + pred = tf.constant([True, False], dtype=tf.bool) + + self._assert_invalid_pred_on_device(data, pred, '/CPU:0') + self._assert_invalid_pred_on_device(data, pred, '/device:MUSA:0') if __name__ == "__main__":