PR tensorflow#21886: [ROCM][NFC] BlasLt interface refactoring & simplifying: part I

Imported from GitHub PR openxla/xla#21886

After this PR tensorflow#73926 is merged, we can remove unnecessary low-level DoMatmul functions from GpuBlasLt interface (which otherwise looks scary and unnecessarily complicated).

Furthermore, we can also remove **ValidateInputs** function from the interface and derived classes since a high-level **ExecuteOnStream** function already handles data-types correctly. This also greatly simplifies the code.

Also, I have packed the input arguments of ExecuteOnStream calls to a struct **MemoryArgs** to simplify arguments passing in derived classes and improve code readability.

Finally, in the original GpuBlasLt PR: openxla/xla#5911, I made a sort of mistake by adding a reference to **blas_lt** to the MatmulPlan class [here](https://github.com/openxla/xla/blob/main/xla/stream_executor/rocm/hip_blas_lt.h#L135), thereby making MatmulPlans bound to a **particular BlasLt instance**. This resulted in some further bugfixes and, most importantly, complicated GpuBlasLt cache design in gpublas_lt_matmul_thunk.cc/.h. In this PR, I remove this reference again from MatmulPlan class and in the next NFC PR the cache mechanics can also be simplified.

Unfortunately, this change also requires a tandem PR for Tensorflow: tensorflow#85835

@xla-rotation Would you please have a look

Copybara import of the project:

--
e96bb2fbedab3f53b31ef0e1748582c76e9fb105 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

blaslt interface refactoring: removing blas_lt_ref

added cuda adaptions

cuda-side adaptions

cuda side adaptions

fix

fixing pointers

Merging this change closes tensorflow#21886

PiperOrigin-RevId: 727898957

Author: pemeliya

tensorflow/core/kernels/matmul_util.cc

@@ -1,8 +1,11 @@
 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
 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.
@@ -176,7 +179,7 @@ StatusOr<se::blas::ComputationType> GetBlasComputationType(
 
     TF_ASSIGN_OR_RETURN(
         auto algorithms,
-        plan->GetAlgorithms(*max_algorithm_count, max_scratch_size));
+        plan->GetAlgorithms(stream, *max_algorithm_count, max_scratch_size));
 
     ptr->second = {std::move(plan), std::move(algorithms)};
   }
@@ -201,9 +204,8 @@ Status PlanAndAlgorithms::ExecuteOnStream(
                                se::DeviceMemoryBase{},  // c_scale_buffer
                                se::DeviceMemoryBase{},  // d_scale_buffer
                                se::DeviceMemoryBase{},  // d_amax_buffer
-                               algorithms[algorithm_idx],
-                               std::nullopt,  // workspace
-                               &scratch_allocator, profile_result);
+                               algorithms[algorithm_idx], scratch_allocator,
+                               profile_result);
 }
 
 }  // namespace tensorflow

third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd.cc

@@ -1205,7 +1205,7 @@ CublasLtCmd::CublasLtCmd(
       workspace_buffer_(workspace_buffer) {}
 
 absl::StatusOr<se::gpu::BlasLt::MatmulPlan*> CublasLtCmd::GetMatmulPlan(
-    const stream_executor::Stream* stream) {
+    const se::Stream* stream) {
   auto it = matmul_plans_cache_.find(stream);
   if (it != matmul_plans_cache_.end()) return it->second.get();
   TF_ASSIGN_OR_RETURN(auto plan, se::gpu::BlasLt::GetMatmulPlan(
@@ -1215,13 +1215,14 @@ absl::StatusOr<se::gpu::BlasLt::MatmulPlan*> CublasLtCmd::GetMatmulPlan(
 }
 
 absl::StatusOr<se::gpu::BlasLt::MatmulAlgorithm>
-CublasLtCmd::GetMatmulAlgorithm(const se::gpu::BlasLt::MatmulPlan* plan,
+CublasLtCmd::GetMatmulAlgorithm(const se::Stream* stream,
+                                const se::gpu::BlasLt::MatmulPlan* plan,
                                 int64_t max_workspace) {
   auto it = matmul_algorithm_cache_.find(plan);
   if (it != matmul_algorithm_cache_.end()) return it->second;
   TF_ASSIGN_OR_RETURN(
       auto algorithms,
-      plan->GetAlgorithms(/*max_algorithm_count*/ 128,
+      plan->GetAlgorithms(stream, /*max_algorithm_count*/ 128,
                           /*max_workspace_size*/ max_workspace));
   TF_RET_CHECK(algorithm_idx_ >= 0 && algorithm_idx_ < algorithms.size());
   auto [it_insert, _] =
@@ -1237,7 +1238,8 @@ absl::Status CublasLtCmd::Initialize(const Thunk::InitializeParams& params,
   // Populate plan and algorithm cache;
   TF_ASSIGN_OR_RETURN(auto plan, GetMatmulPlan(params.stream));
   TF_RETURN_IF_ERROR(
-      GetMatmulAlgorithm(plan, workspace_buffer_.size()).status());
+      GetMatmulAlgorithm(params.stream, plan, workspace_buffer_.size())
+          .status());
   return absl::OkStatus();
 }
 
@@ -1246,7 +1248,8 @@ absl::Status CublasLtCmd::Record(const Thunk::ExecuteParams& execute_params,
                                  se::CommandBuffer* command_buffer) {
   TF_ASSIGN_OR_RETURN(auto plan, GetMatmulPlan(execute_params.stream));
   TF_ASSIGN_OR_RETURN(auto algorithm,
-                      GetMatmulAlgorithm(plan, workspace_buffer_.size()));
+                      GetMatmulAlgorithm(execute_params.stream, plan,
+                                         workspace_buffer_.size()));
 
   const BufferAllocations& allocs = *execute_params.buffer_allocations;
 

third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd.h

@@ -803,13 +803,13 @@ class CublasLtCmd : public TracedCommandBufferCmd {
 
  private:
   absl::StatusOr<se::gpu::BlasLt::MatmulPlan*> GetMatmulPlan(
-      const stream_executor::Stream* stream);
+      const se::Stream* stream);
 
   absl::StatusOr<se::gpu::BlasLt::MatmulAlgorithm> GetMatmulAlgorithm(
-      const se::gpu::BlasLt::MatmulPlan* plan, int64_t max_workspace);
+      const se::Stream* stream, const se::gpu::BlasLt::MatmulPlan* plan,
+      int64_t max_workspace);
 
-  absl::flat_hash_map<const stream_executor::Stream*,
-                      se::gpu::BlasLt::MatmulPlanPtr>
+  absl::flat_hash_map<const se::Stream*, se::gpu::BlasLt::MatmulPlanPtr>
       matmul_plans_cache_;
 
   absl::flat_hash_map<const se::gpu::BlasLt::MatmulPlan*,

third_party/xla/xla/backends/gpu/runtime/gpublas_lt_matmul_thunk.cc

@@ -65,11 +65,11 @@ CublasLtMatmulThunk::CublasLtMatmulThunk(
 absl::Status CublasLtMatmulThunk::ExecuteOnStream(const ExecuteParams& params) {
   TF_ASSIGN_OR_RETURN(auto plan, GetMatmulPlan(params.stream));
 
-  TF_ASSIGN_OR_RETURN(
-      auto algorithm,
-      GetMatmulAlgorithm(plan, workspace_buffer_.has_value()
-                                   ? workspace_buffer_.value().size()
-                                   : 0));
+  TF_ASSIGN_OR_RETURN(auto algorithm,
+                      GetMatmulAlgorithm(params.stream, plan,
+                                         workspace_buffer_.has_value()
+                                             ? workspace_buffer_.value().size()
+                                             : 0));
 
   VLOG(3) << "Running cublas_lt matmul thunk";
   const BufferAllocations& allocs = *params.buffer_allocations;
@@ -99,7 +99,7 @@ absl::Status CublasLtMatmulThunk::ExecuteOnStream(const ExecuteParams& params) {
     aux = allocs.GetDeviceAddress(aux_buffer_);
   }
 
-  std::optional<se::DeviceMemoryBase> workspace;
+  se::DeviceMemoryBase workspace;
   if (workspace_buffer_.has_value()) {
     workspace = allocs.GetDeviceAddress(workspace_buffer_.value());
   }
@@ -112,7 +112,7 @@ absl::Status CublasLtMatmulThunk::ExecuteOnStream(const ExecuteParams& params) {
 }
 
 absl::StatusOr<se::gpu::BlasLt::MatmulPlan*> CublasLtMatmulThunk::GetMatmulPlan(
-    const stream_executor::Stream* stream) {
+    const se::Stream* stream) {
   {
     absl::MutexLock lock(&matmul_plans_cache_mutex_);
     auto it = matmul_plans_cache_.find(stream);
@@ -127,7 +127,8 @@ absl::StatusOr<se::gpu::BlasLt::MatmulPlan*> CublasLtMatmulThunk::GetMatmulPlan(
 }
 
 absl::StatusOr<se::gpu::BlasLt::MatmulAlgorithm>
-CublasLtMatmulThunk::GetMatmulAlgorithm(const se::gpu::BlasLt::MatmulPlan* plan,
+CublasLtMatmulThunk::GetMatmulAlgorithm(const se::Stream* stream,
+                                        const se::gpu::BlasLt::MatmulPlan* plan,
                                         int64_t max_workspace) {
   {
     absl::MutexLock lock(&matmul_algorithm_cache_mutex_);
@@ -136,7 +137,8 @@ CublasLtMatmulThunk::GetMatmulAlgorithm(const se::gpu::BlasLt::MatmulPlan* plan,
   }
   TF_ASSIGN_OR_RETURN(
       auto algorithms,
-      plan->GetAlgorithms(/*max_algorithm_count*/ 128,
+      plan->GetAlgorithms(stream,
+                          /*max_algorithm_count*/ 128,
                           /*max_workspace_size*/ max_workspace));
   TF_RET_CHECK(algorithm_idx_ >= 0 && algorithm_idx_ < algorithms.size());
 

third_party/xla/xla/backends/gpu/runtime/gpublas_lt_matmul_thunk.h

@@ -74,7 +74,8 @@ class CublasLtMatmulThunk : public Thunk {
   absl::StatusOr<se::gpu::BlasLt::MatmulPlan*> GetMatmulPlan(
       const stream_executor::Stream* stream);
   absl::StatusOr<se::gpu::BlasLt::MatmulAlgorithm> GetMatmulAlgorithm(
-      const se::gpu::BlasLt::MatmulPlan* plan, int64_t max_workspace);
+      const se::Stream* stream, const se::gpu::BlasLt::MatmulPlan* plan,
+      int64_t max_workspace);
 
   absl::Mutex matmul_plans_cache_mutex_;
   absl::flat_hash_map<const stream_executor::Stream*,

third_party/xla/xla/service/gpu/autotuning/gemm_algorithm_picker.cc

@@ -186,7 +186,7 @@ class GemmAutotuner {
 
     TF_ASSIGN_OR_RETURN(
         auto algorithms,
-        plan->GetAlgorithms(/*max_algorithm_count*/ 128,
+        plan->GetAlgorithms(stream_, /*max_algorithm_count*/ 128,
                             /*max_workspace_size*/ workspace_buffer.size()));
 
     auto tuned_func = [&](const BlasLt::MatmulAlgorithm& algorithm)