Feat: support for XLA

CMakeLists.txt

@@ -158,6 +158,11 @@ endif()
 
 # Link paths
 link_directories("${MUSA_PATH}/lib" "${MUSA_PATH}/lib64" "${TF_LIB_DIR}")
+set(TF_CC_LIBRARY "")
+if(EXISTS "${TF_LIB_DIR}/libtensorflow_cc.so.2")
+    set(TF_CC_LIBRARY "${TF_LIB_DIR}/libtensorflow_cc.so.2")
+    message(STATUS "TensorFlow CC Library: ${TF_CC_LIBRARY}")
+endif()
 
 # Automatically collect source files
 file(GLOB_RECURSE CPP_SOURCES
@@ -221,7 +226,7 @@ set_target_properties(musa_plugin PROPERTIES
 )
 
 # Link
-target_link_libraries(musa_plugin PRIVATE ${TF_LINK_FLAGS} musa mublas ${MUSA_DNN_LIBRARY} rt)
+target_link_libraries(musa_plugin PRIVATE ${TF_LINK_FLAGS} ${TF_CC_LIBRARY} musa mublas ${MUSA_DNN_LIBRARY} rt)
 
 # Generate Python runtime configuration bindings.
 add_library(_runtime_config_bindings MODULE

musa_ext/kernels/fusion/musa_concat_matmul_op.cc

@@ -159,6 +159,9 @@ REGISTER_OP("MusaConcatMatMul")
     .Attr("transpose_b: bool = false")
     .Attr("num_concat: int >= 1")
     .Attr("concat_input_idx: int")
-    .SetShapeFn(shape_inference::MatMulShape);
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      c->set_output(0, c->UnknownShape());
+      return OkStatus();
+    });
 
 }  // namespace tensorflow

musa_ext/kernels/fusion/xla/musa_clip_prelu_xla.cc

@@ -0,0 +1,51 @@
+#include "musa_fusion_xla_common.h"
+
+namespace tensorflow {
+namespace {
+
+class MusaClipXlaOp : public XlaOpKernel {
+ public:
+  explicit MusaClipXlaOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    BCast bcast_x_lo(BCast::FromShape(ctx->InputShape(0)),
+                     BCast::FromShape(ctx->InputShape(1)));
+    OP_REQUIRES(ctx, bcast_x_lo.IsValid(),
+                errors::InvalidArgument("MusaClip x/lo shapes are not "
+                                        "broadcast-compatible"));
+    BCast bcast_all(bcast_x_lo.output_shape(),
+                    BCast::FromShape(ctx->InputShape(2)));
+    OP_REQUIRES(ctx, bcast_all.IsValid(),
+                errors::InvalidArgument("MusaClip x/lo/hi shapes are not "
+                                        "broadcast-compatible"));
+
+    TensorShape out_shape = BCast::ToShape(bcast_all.output_shape());
+    xla::XlaOp x = ctx->Input(0);
+    xla::XlaOp lo = ctx->Input(1);
+    xla::XlaOp hi = ctx->Input(2);
+    BroadcastToShape(ctx, &x, out_shape);
+    BroadcastToShape(ctx, &lo, out_shape);
+    BroadcastToShape(ctx, &hi, out_shape);
+    ctx->SetOutput(0, xla::Clamp(lo, x, hi));
+  }
+};
+
+class MusaPReluXlaOp : public XlaOpKernel {
+ public:
+  explicit MusaPReluXlaOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp x = ctx->Input(0);
+    xla::XlaOp alpha = ctx->Input(1);
+    BroadcastToShape(ctx, &alpha, ctx->InputShape(0));
+    xla::XlaOp zero = xla::ZerosLike(x);
+    ctx->SetOutput(0, xla::Max(zero, x) + alpha * xla::Min(zero, x));
+  }
+};
+
+REGISTER_XLA_OP(Name("MusaClip"), MusaClipXlaOp);
+REGISTER_XLA_OP(Name("MusaPRelu").TypeConstraint("T", kFloatTypes),
+                MusaPReluXlaOp);
+
+}  // namespace
+}  // namespace tensorflow