PR #10638: Handle no-op custom call in the emitter

Imported from GitHub PR openxla/xla#10638

The AllocateBuffer custom call (a.k.a. `kNopCustomCallTarget`) is no-op because the runtime allocates the buffer. Let us handle AllocateBuffer by emitting nothing and returning.
Copybara import of the project:

--
6ad7c399e05182a543e50ed14f34b99a281155ce by Jaroslav Sevcik <jsevcik@nvidia.com>:

Emit no-op (AllocateBuffer custom-call)

--
0f78007475921031549c8555ad56ed8903efe6cb by Jaroslav Sevcik <jsevcik@nvidia.com>:

Renamed allocate-buffer to nop in the test

Merging this change closes #10638

PiperOrigin-RevId: 628129950

third_party/xla/xla/service/gpu/ir_emitter_unnested.cc

@@ -2921,6 +2921,9 @@ absl::Status IrEmitterUnnested::EmitHloInstruction(
       if (instr->custom_call_target() == "__gpu$xla.gpu.triton") {
         return EmitTritonCustomCall(custom_call);
       }
+      if (instr->custom_call_target() == kNopCustomCallTarget) {
+        return absl::OkStatus();
+      }
       return EmitCustomCallThunk(custom_call);
     }
     case HloOpcode::kFusion: {

third_party/xla/xla/service/gpu/tests/BUILD

@@ -1080,3 +1080,15 @@ xla_cc_test(
         "@local_tsl//tsl/platform:test_main",
     ],
 )
+
+xla_cc_test(
+    name = "nop_custom_call_test",
+    srcs = ["nop_custom_call_test.cc"],
+    tags = tf_cuda_tests_tags(),
+    deps = [
+        "//xla:xla_proto_cc",
+        "//xla/service:gpu_plugin",
+        "//xla/tests:hlo_test_base",
+        "@local_tsl//tsl/platform:test_main",
+    ],
+)

third_party/xla/xla/service/gpu/tests/nop_custom_call_test.cc

@@ -0,0 +1,51 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/tests/hlo_test_base.h"
+
+namespace xla {
+namespace gpu {
+namespace {
+
+class NopCustomCallTest : public HloTestBase {};
+
+TEST_F(NopCustomCallTest, RunAllocateBufferAndUpdate) {
+  // The test uses a custom call with the AllocateBuffer target (also known as
+  // kNopCustomCallTarget) to allocate an output buffer. Then it verifies
+  // we can successfully modify the buffer.
+  const char* hlo_text = R"(
+  HloModule AllocateBuffer, is_scheduled=true
+
+  overwrite_one {
+    p0 = s32[1] parameter(0)
+    c0 = s32[] constant(0)
+    c1 = s32[1] constant({1})
+    ROOT dus0 = s32[1] dynamic-update-slice(p0, c1, c0)
+  }
+
+  ENTRY main {
+    buffer = s32[1] custom-call(), custom_call_target="AllocateBuffer"
+    ROOT fusion = s32[1] fusion(buffer), kind=kLoop, calls=overwrite_one
+  })";
+  auto module = ParseAndReturnVerifiedModule(hlo_text).value();
+
+  Literal result = ExecuteNoHloPasses(std::move(module), {});
+  Literal expected = LiteralUtil::CreateR1<int32_t>({1});
+  EXPECT_TRUE(LiteralTestUtil::Equal(expected, result));
+}
+
+}  // namespace
+}  // namespace gpu
+}  // namespace xla