[XLA:GPU][NFC] Rename gpu_version_ to compute_capability_ in `gem…

…m_fusion`. PiperOrigin-RevId: 647695512
tensorflow · Jun 28, 2024 · f346a1f · f346a1f
1 parent 32d86c6
commit f346a1f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/third_party/xla/xla/service/gpu/gemm_fusion.cc b/third_party/xla/xla/service/gpu/gemm_fusion.cc
@@ -798,13 +798,14 @@ bool ShouldTritonHandleGEMM(HloDotInstruction& dot,
 absl::StatusOr<bool> GemmFusion::Run(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
-  TF_RETURN_IF_ERROR(EnsureTritonSupportsComputeCapability(gpu_version_));
+  TF_RETURN_IF_ERROR(
+      EnsureTritonSupportsComputeCapability(compute_capability_));
 
   bool changed = false;
   for (HloComputation* computation :
        module->MakeNonfusionComputations(execution_threads)) {
     TF_ASSIGN_OR_RETURN(bool result,
-                        RunOnComputation(computation, gpu_version_));
+                        RunOnComputation(computation, compute_capability_));
     changed |= result;
   }
   return changed;

diff --git a/third_party/xla/xla/service/gpu/gemm_fusion.h b/third_party/xla/xla/service/gpu/gemm_fusion.h
@@ -38,8 +38,8 @@ bool ShouldTritonHandleGEMM(HloDotInstruction&,
 // that target Triton-based matmul emitter.
 class GemmFusion : public HloModulePass {
  public:
-  explicit GemmFusion(const se::GpuComputeCapability& gpu_version)
-      : gpu_version_(gpu_version) {}
+  explicit GemmFusion(const se::GpuComputeCapability& compute_capability)
+      : compute_capability_(compute_capability) {}
   absl::string_view name() const override { return "triton-gemm-rewriter"; }
 
   using HloPassInterface::Run;
@@ -48,7 +48,7 @@ class GemmFusion : public HloModulePass {
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
  private:
-  se::GpuComputeCapability gpu_version_;
+  se::GpuComputeCapability compute_capability_;
 };
 
 }  // namespace gpu