[Bugfix] Fix safe memory legalization for fragment store (#446)

* [Enhancement] Improve layout inference accuracy in ParallelOp (#441)

* Added logic to use non-replicated buffers as source buffers for more accurate layout inference.
* Enhanced comments to clarify the rationale behind buffer selection in layout inference process.

* [Enhancement] Add error handling macros and refactor loop partitioning logic

* Introduced TILELANG_CHECK macro for improved error handling in CUDA and HIP code, providing detailed error messages for kernel launches.
* Enhanced loop partitioning logic to handle fragment buffers more effectively, ensuring correct replication based on thread extent.
* Added logging for thread range in PlanLoopPartition to aid in debugging and performance analysis.
* Updated pass configuration management to streamline vectorization control in the optimization process.

* lint fix

* remove debug print

* [Refactor] Update legalize_safe_memory_access.cc to improve memory access handling

* Replaced Apache License header with MIT License.
* Added logic to handle local buffer conditions in memory access.
* Introduced IsLocalBuffer function to check buffer scope.
* Enhanced comments for clarity on memory access operations.

src/transform/legalize_safe_memory_access.cc

-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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.
- */
-
 /*!
- * \file layout_inference.cc
- * \brief infer the fragment/shared memory layout
+ * \file legalize_safe_memory_access.cc
+ * \brief legalize safe memory access
  */
 
 #include <tvm/tir/builtin.h>
@@ -161,7 +142,6 @@ private:
     GlobalMemChecker checker(analyzer_);
     checker(store);
     Array<PrimExpr> conditions = checker.GetConditions();
-
     if (conditions.size() == 0) {
       return store;
     }
@@ -182,6 +162,18 @@ private:
       }
       store.CopyOnWrite()->value = value;
       return store;
+    } else if (IsLocalBuffer(store->buffer)) {
+      PrimExpr value = store->value;
+      for (auto cond : conditions) {
+        ICHECK(cond.dtype() == DataType::Bool(1))
+            << "condition is not a boolean: " << cond;
+        value = if_then_else(cond, value, make_zero(value->dtype));
+      }
+      store.CopyOnWrite()->value = value;
+      return store;
+    } else {
+      LOG(FATAL) << "Check store buffer: " << store->buffer
+                 << " is not a global or shared or local buffer";
     }
 
     return store;
@@ -215,6 +207,11 @@ private:
     return evaluate;
   }
 
+  bool IsLocalBuffer(const Buffer &buffer) {
+    String scope = buffer.scope();
+    return scope == "local" || scope == "local.fragment";
+  }
+
   bool isSharedBuffer(const Buffer &buffer) {
     String scope = buffer.scope();
     return scope == "shared" || scope == "shared.dyn";