Skip to content

Albert/fix wrong op define #284

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
real-zhangzhe merged 2 commits into from
May 27, 2026
Merged

Albert/fix wrong op define #284

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c97181a
albert/fix_wrong_op_define
real-zhangzhe May 27, 2026
97e1898
albert/fix_wrong_op_define
real-zhangzhe May 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions musa_ext/graph/fusion/tensordot_fusion.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -821,6 +821,40 @@ Status MusaTensorDotFusion::Apply(GraphDef* graph,
fuse_node_names.erase(name);
}

// 递归保留 shared node 的上游依赖；否则虽然 shared node 本身会保留，
// 但它依赖的内部 shape/perm 辅助节点仍可能被删掉，留下悬空输入。
std::vector<std::string> keep_worklist(shared_nodes.begin(),
shared_nodes.end());
std::unordered_set<std::string> 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";

Expand Down
6 changes: 3 additions & 3 deletions musa_ext/kernels/array/musa_matrix_set_diag_kernel.mu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -112,7 +112,7 @@ void MusaMatrixSetDiagKernelLauncher(
musaStream_t, typename TTypes<Scalar, 3>::ConstTensor&, typename TTypes<Scalar, 3>::ConstTensor&, \
typename TTypes<Scalar, 3>::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);
Expand All @@ -122,4 +122,4 @@ INSTANTIATE_MATRIX_SET_DIAG_KERNEL(Eigen::bfloat16);
INSTANTIATE_MATRIX_SET_DIAG_KERNEL(bool);

} // namespace musa
} // namespace tensorflow
} // namespace tensorflow
24 changes: 12 additions & 12 deletions musa_ext/kernels/array/musa_matrix_set_diag_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename T>
Expand Down Expand Up @@ -159,7 +159,7 @@ class MusaMatrixSetDiagOp : public MusaOpKernel {
musaStream_t stream = GetMusaStreamByCtx(context);
MusaMatrixSetDiagKernelLauncher<T>(
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_);
}

Expand All @@ -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<type>("T"), \
MusaMatrixSetDiagOp<type>); \
REGISTER_KERNEL_BUILDER( \
Name("MusaMatrixSetDiagV2").Device("MUSA").TypeConstraint<type>("T"), \
MusaMatrixSetDiagOp<type>); \
REGISTER_KERNEL_BUILDER( \
Name("MusaMatrixSetDiagV3").Device("MUSA").TypeConstraint<type>("T"), \
#define REGISTER_MATRIX_SET_DIAG_KERNEL(type) \
REGISTER_KERNEL_BUILDER( \
Name("MatrixSetDiag").Device("MUSA").TypeConstraint<type>("T"), \
MusaMatrixSetDiagOp<type>); \
REGISTER_KERNEL_BUILDER( \
Name("MatrixSetDiagV2").Device("MUSA").TypeConstraint<type>("T"), \
MusaMatrixSetDiagOp<type>); \
REGISTER_KERNEL_BUILDER( \
Name("MatrixSetDiagV3").Device("MUSA").TypeConstraint<type>("T"), \
MusaMatrixSetDiagOp<type>);

REGISTER_MATRIX_SET_DIAG_KERNEL(float);
Expand All @@ -192,4 +192,4 @@ REGISTER_MATRIX_SET_DIAG_KERNEL(bool);

#undef REGISTER_MATRIX_SET_DIAG_KERNEL
} // namespace musa
} // namespace tensorflow
} // namespace tensorflow
8 changes: 4 additions & 4 deletions musa_ext/kernels/control_flow/musa_switch_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,9 @@ class MusaSwitchOp : public MusaOpKernel {
}
}; // class MusaSwitchOp

#define REGISTER_MUSA_SWITCH(type) \
REGISTER_KERNEL_BUILDER( \
Name("MusaSwitch").Device(DEVICE_MTGPU).TypeConstraint<type>("T"), \
#define REGISTER_MUSA_SWITCH(type) \
REGISTER_KERNEL_BUILDER( \
Name("Switch").Device(DEVICE_MTGPU).TypeConstraint<type>("T"), \
MusaSwitchOp<type>);

REGISTER_MUSA_SWITCH(float);
Expand All @@ -43,4 +43,4 @@ REGISTER_MUSA_SWITCH(uint8);
REGISTER_MUSA_SWITCH(bool);

} // namespace musa
} // namespace tensorflow
} // namespace tensorflow
7 changes: 7 additions & 0 deletions musa_ext/kernels/fusion/musa_sigmoid_calibration_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename T>
void LaunchSigmoidCalibrationKernel(const void*, const void*, void*, int,
musaStream_t);
Expand Down
46 changes: 30 additions & 16 deletions test/ops/switch_op_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]:
Expand All @@ -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__":
Expand Down
Loading